# -*- coding: utf-8 -*-
"""
从第1行到第739行这部分python代码，对应的是matlab的第1行到第739行
对应Matlab文件: Func_VPE.m
功能: 向量抛物方程(VPE)波导/隧道内传播求解器

说明:
- 为尽量逐行对应，本实现保持Matlab变量名和流程；
- 索引差异: Matlab为1基索引，Python/numpy为0基索引。为最大限度避免逻辑偏差，
  本文件中的二维/三维数组统一申请为“多一行多一列”的尺寸，并在访问时直接使用Matlab风格的索引值，
  即: 若Matlab访问 A(i,j)，此处直接使用 A[i, j] (i,j均为Matlab计算得到的索引)。因此数组第0行/第0列未使用。
- 切片差异: Matlab切片是“闭区间”，Python切片是“左闭右开”。遇到Matlab a:b 时，转为 Python a:b+1。
- 三角函数差异: Matlab使用度数版本 atand/sind/cosd，这里提供等价的度数版本函数。
"""

from typing import Tuple
import numpy as np

import matplotlib.pyplot as plt
import time

class TimeCountVar:
    startTime = time.time()

def timeCount(msg):
    endTime = time.time()
    # calculate the milliseconds:
    ms = (endTime - TimeCountVar.startTime) * 1000
    print(f"{msg}: {ms:.2f} ms")
    TimeCountVar.startTime = endTime

# 从第740行到第749行这部分python代码，对应的是matlab的第740行到第749行
# Matlab自定义的V_roundn: 这里统一等价为按8位小数四舍五入
def v_roundn(x):
    x = np.asarray(x)
    return np.round(x, 8)


# Matlab的atand/sind/cosd等价实现
def atand(x):
    # x = np.asarray(x)
    # 处理正负无穷
    # 对于标量inf/-inf保留为±90
    if np.isscalar(x):
        if np.isposinf(x):
            return 90.0
        if np.isneginf(x):
            return -90.0
        return float(np.degrees(np.arctan(x)))
    # 对于数组，位置替换
    res = np.degrees(np.arctan(x))
    res = np.asarray(res)
    res = res.astype(float)
    res[np.isposinf(x)] = 90.0
    res[np.isneginf(x)] = -90.0
    return res


def sind(theta):
    return np.sin(np.deg2rad(theta))


def cosd(theta):
    return np.cos(np.deg2rad(theta))


# 工具: 在矩阵对角线上累加向量(vec)，偏移k (k>0为上对角线, k<0为下对角线)
def add_diag(M: np.ndarray, vec: np.ndarray, k: int) -> None:
    n = M.shape[0]
    vec = np.asarray(vec)
    if k >= 0:
        i = np.arange(0, n - k)
        M[i, i + k] += vec[: i.shape[0]]
    else:
        k2 = -k
        i = np.arange(0, n - k2)
        M[i + k2, i] += vec[: i.shape[0]]


def func_vpe(
    f: float,
    eps_r: float,
    sigma: float,
    distance: float,
    TX: list,
    k_w0: float,
    shape: int,
    BC: int,
    cross_section: list,
    dx: float,
    ds: float,
) -> Tuple[np.ndarray, np.ndarray]:
    """
    从第1行到第739行这部分python代码，对应的是matlab的第1行到第739行
    Matlab: function [Ex_all, Ey_all] = Func_VPE(f, eps_r, sigma, distance, TX, k_w0, shape, BC, cross_section, dx, ds)
    返回: Ex_all, Ey_all

    重要索引说明:
    - 为匹配Matlab的1基索引，本函数中所有工作数组均在第0行/第0列留空，从1开始使用。
    - 例如Matlab的大小(Ny+1, Nx+1)在Python中申请为同样(Ny+1+1, Nx+1+1)= (Ny+2, Nx+2)，并只使用[1..Ny+1, 1..Nx+1]区域。
    """
    # 从第43行到第56行，对应matlab的第43行到第56行 (输入参数与常量)
    
    TimeCountVar.startTime = time.time()
    c = 299792458.0
    lambda_ = c / f
    k0 = 2 * np.pi / lambda_
    mu_0 = 4 * np.pi * 1e-7
    eps_0 = 1 / mu_0 / c / c
    sigma_r = sigma / (2 * np.pi * f * eps_0)
    eps_rc = eps_r - 1j * sigma_r
    Z = np.sqrt(eps_rc - 1) / eps_rc
    s_start = 0.0
    s_end = distance

    # 从第58行到第90行，对应matlab的第58行到第90行 (几何参数按shape分类)
    if shape == 1:
        H_top = cross_section[1] / 2
        W_half = cross_section[0] / 2
        H_bottom = H_top
        R2_bottom = 0
        shape_top = 4
        shape_bottom = 4
    elif shape == 2:
        H_top = cross_section[1]
        W_half = cross_section[0]
        H_bottom = cross_section[3]
        R2_bottom = cross_section[2]
        shape_top = 1
        shape_bottom = 1
    elif shape == 3:
        H_top = cross_section[1]
        W_half = cross_section[0]
        H_bottom = cross_section[3]
        R2_bottom = cross_section[2]
        shape_top = 1
        shape_bottom = 4
    elif shape == 4:
        H_top = cross_section[1]
        W_half = cross_section[0]
        H_bottom = cross_section[3]
        R2_bottom = cross_section[2]
        x_B1 = cross_section[4]
        y_B2 = cross_section[5]
        shape_top = 5
        shape_bottom = 5
    else:
        raise ValueError("Unsupported shape: {}".format(shape))

    # 从第93行到第110行，对应matlab的第93行到第110行 (网格离散)
    dy = dx
    nh_top = int(np.ceil(v_roundn(H_top / dy)))
    nh_bottom = int(np.ceil(v_roundn(H_bottom / dy)))
    nw_half = int(np.ceil(v_roundn(W_half / dx)))
    Nx = 2 * nw_half
    Ny = nh_top + nh_bottom
    x_c = int(Nx / 2)   # Matlab中心列索引 (1基)
    y_c = int(nh_top)   # Matlab中心行索引 (1基)
    Nz = int(np.ceil((s_end - s_start) / ds))

    # 初始化交界面几何辅助变量 (多形状分支会赋值)
    Hs_top = 0.0
    X_t = 0.0
    R1_top = 0.0
    R2_top = 0.0
    p1_top = 0.0
    p2_top = 0.0
    t1_top = 0.0
    t2_top = 0.0
    x_B_t = None
    y_B_t = None
    n_points_t = None
    n_seg_t = None

    # 从第112行到第134行，对应matlab的第112行到第134行 (顶部边界形状参数)
    if shape_top == 1:
        Hs_top = 0.0
        X_t = 0.0
        R2_top = H_top + Hs_top
        R1_top = np.sqrt(W_half ** 2 / (1 - (Hs_top / R2_top) ** 2))
    elif shape_top == 2:
        X_t = 0.0
        p2_top = H_top
        p1_top = -p2_top / (W_half * W_half)
    elif shape_top == 3:
        X_t = x_B1
        t1_top = H_top / (X_t - W_half)
        t2_top = -t1_top * W_half
    elif shape_top == 4:
        pass
    elif shape_top == 5:
        x_B_t = np.array([0.0, x_B1, W_half, W_half])
        y_B_t = np.array([H_top, H_top, y_B2 - H_bottom, 0.0])
        n_points_t = len(x_B_t)
        n_seg_t = n_points_t - 1

    # 从第136行到第171行，对应matlab的第136行到第171行 (底部边界形状参数)
    Hs_bottom = 0.0
    X_b = 0.0
    R1_bottom = 0.0
    p1_bottom = 0.0
    p2_bottom = 0.0
    t1_bottom = 0.0
    t2_bottom = 0.0
    x_B_b = None
    y_B_b = None
    n_points_b = None
    n_seg_b = None
    if shape_bottom == 1:
        Hs_bottom = 0.0
        X_b = np.sqrt(W_half ** 2 - H_bottom ** 2)
        R1_bottom = np.sqrt(
            W_half ** 2 + Hs_bottom ** 2 * (W_half ** 2 - X_b ** 2) / (H_bottom ** 2 + 2 * H_bottom * Hs_bottom)
        )
        R2_bottom = np.sqrt(
            ((H_bottom ** 2 + 2 * H_bottom * Hs_bottom) * R1_bottom ** 2) / (W_half ** 2 - X_b ** 2)
        )
    elif shape_bottom == 2:
        X_b = 1.75
        p1_bottom = -H_bottom / (X_b ** 2 - W_half ** 2)
        p2_bottom = -p1_bottom * W_half ** 2
    elif shape_bottom == 3:
        X_b = x_B1
        t1_bottom = -H_bottom / (X_b - W_half)
        t2_bottom = -t1_bottom * W_half
    elif shape_bottom == 4:
        pass
    elif shape_bottom == 5:
        x_B_b = np.array([W_half, W_half, 0.0])
        y_B_b = -np.array([0.0, H_bottom, H_bottom])
        n_points_b = len(x_B_b)
        n_seg_b = n_points_b - 1

    # 从第176行到第193行，对应matlab的第176行到第193行 (系数矩阵初始化)
    a_x = np.full((Ny + 1, Nx + 1), np.inf, dtype=float)
    R_x1 = np.zeros((Ny + 1, Nx + 1), dtype=float)
    R_x2 = np.zeros((Ny + 1, Nx + 1), dtype=float)
    R_x3 = np.zeros((Ny + 1, Nx + 1), dtype=complex)
    R_x2_c = np.zeros((Ny + 1, Nx + 1), dtype=float)
    R_x2_n = np.zeros((Ny + 1, Nx + 1), dtype=float)
    R_in_d = np.zeros((Ny + 1, Nx + 1), dtype=float)
    Z_x1 = np.zeros((Ny + 1, Nx + 1), dtype=complex)
    Z_x2 = np.zeros((Ny + 1, Nx + 1), dtype=complex)
    Z_x3 = np.zeros((Ny + 1, Nx + 1), dtype=complex)
    Z_x4 = np.zeros((Ny + 1, Nx + 1), dtype=complex)

    # 从第195行到第217行，对应matlab的第195行到第217行 (顶部几何xx_t)
    yy_t = np.arange(nh_top - 1, -1, -1) * dy
    if shape_top == 1:
        xx_t = R1_top * np.sqrt(1 - (yy_t + Hs_top) ** 2 / R2_top ** 2)
    elif shape_top == 2:
        xx_t = np.sqrt((yy_t - p2_top) / p1_top)
    elif shape_top == 3:
        xx_t = (yy_t - t2_top) / t1_top
    elif shape_top == 4:
        xx_t = W_half * np.ones_like(yy_t)
    elif shape_top == 5:
        k_B_t = np.zeros(n_seg_t)
        c_B_t = np.zeros(n_seg_t)
        xx_t_list = []
        for i in range(n_seg_t):
            k_B_t[i] = (y_B_t[i + 1] - y_B_t[i]) / (x_B_t[i + 1] - x_B_t[i])
            c_B_t[i] = y_B_t[i] - k_B_t[i] * x_B_t[i]
            y_start = int(np.ceil(v_roundn(y_B_t[i] / dy)))
            y_end = int(np.ceil(v_roundn(y_B_t[i + 1] / dy)))
            yy_t_tentative = np.arange(y_start - 1, y_end - 1, -1) * dy
            xx_t_tentative = (yy_t_tentative - c_B_t[i]) / k_B_t[i]
            xx_t_list.append(xx_t_tentative)
        xx_t = np.concatenate(xx_t_list) if len(xx_t_list) > 0 else np.array([], dtype=float)
        xx_t = np.where(np.isnan(xx_t), x_B_t[-1], xx_t)

    # 从第219行到第243行，对应matlab的第219行到第243行 (底部几何xx_b)
    yy_b = -(np.arange(1, nh_bottom, 1)) * dy
    if shape_bottom == 1:
        xx_b = R1_bottom * np.sqrt(1 - (yy_b - Hs_bottom) ** 2 / R2_bottom ** 2)
    elif shape_bottom == 2:
        xx_b = np.sqrt((yy_b - p2_bottom) / p1_bottom)
    elif shape_bottom == 3:
        xx_b = (yy_b - t2_bottom) / t1_bottom
    elif shape_bottom == 4:
        xx_b = W_half * np.ones_like(yy_b)
    elif shape_bottom == 5:
        k_B_b = np.zeros(n_seg_b)
        c_B_b = np.zeros(n_seg_b)
        xx_b_list = []
        i = 0
        k_B_b[i] = (y_B_b[i + 1] - y_B_b[i]) / (x_B_b[i + 1] - x_B_b[i])
        c_B_b[i] = y_B_b[i] - k_B_b[i] * x_B_b[i]
        y_end = int(np.floor(v_roundn(abs(y_B_b[i + 1]) / dy)))
        yy_b_tentative = -(np.arange(1, y_end + 1) * dy)
        xx_b_tentative = (yy_b_tentative - c_B_b[i]) / k_B_b[i]
        xx_b_list.append(xx_b_tentative)
        for i in range(1, n_seg_b):
            k_B_b[i] = (y_B_b[i + 1] - y_B_b[i]) / (x_B_b[i + 1] - x_B_b[i])
            c_B_b[i] = y_B_b[i] - k_B_b[i] * x_B_b[i]
            y_start = int(np.floor(v_roundn(abs(y_B_b[i]) / dy)))
            y_end = int(np.floor(v_roundn(abs(y_B_b[i + 1]) / dy)))
            yy_b_tentative = -(np.arange(y_start + 1, y_end + 1) * dy)
            xx_b_tentative = (yy_b_tentative - c_B_b[i]) / k_B_b[i]
            xx_b_list.append(xx_b_tentative)
        xx_b = np.concatenate(xx_b_list) if len(xx_b_list) > 0 else np.array([], dtype=float)
        xx_b = np.where(np.isnan(xx_b), x_B_b[0], xx_b)

    # 使用已有Python版 Func_Loc_Res_Coef: 这里内联实现以减少依赖
    from func_loc_res_coef import func_loc_res_coef

    Ix_t, Rx_t = func_loc_res_coef(v_roundn(xx_t / dx))
    Ix_b, Rx_b = func_loc_res_coef(v_roundn(xx_b / dx))

    # set mark for time counting:
    
    # p1
    i_idx = 0
    for p in range(nh_top - 1, -1, -1):
        # 计算theta
        if shape_top == 1:
            theta = atand((R1_top * (p * dy + Hs_top)) / (R2_top * np.sqrt(R2_top ** 2 - (p * dy + Hs_top) ** 2)))
        elif shape_top == 2:
            theta = atand(-1 / 2 / np.sqrt(p1_top * (p * dy - p2_top)))
        elif shape_top == 3:
            if p == 0:
                if shape_bottom == 1:
                    theta_b = atand((R1_bottom * (p * dy - Hs_bottom)) / (R2_bottom * np.sqrt(R2_bottom ** 2 - (p * dy - Hs_bottom) ** 2)))
                elif shape_bottom == 2:
                    theta_b = atand(-1 / 2 / np.sqrt(p1_bottom * (p * dy - p2_bottom)))
                elif shape_bottom == 3:
                    theta_b = atand(-1 / t1_bottom)
                elif shape_bottom == 4:
                    theta_b = atand(0)
                elif shape_bottom == 5:
                    k_B_b_tem = (y_B_b[1] - y_B_b[0]) / (x_B_b[1] - x_B_b[0])
                    theta_b = atand(-1 / k_B_b_tem)
                theta = (atand(-1 / t1_top) + theta_b) / 2
            else:
                theta = atand(-1 / t1_top)
        elif shape_top == 4:
            theta = atand(0)
        elif shape_top == 5:
            i_index = np.where(y_B_t > p * dy)[0]
            i_index = i_index[-1] if i_index.size > 0 else 0
            if y_B_t[i_index + 1] == p * dy and y_B_t[i_index + 1] == 0:
                if shape_bottom == 1:
                    theta_b = atand((R1_bottom * (p * dy - Hs_bottom)) / (R2_bottom * np.sqrt(R2_bottom ** 2 - (p * dy - Hs_bottom) ** 2)))
                elif shape_bottom == 2:
                    theta_b = atand(-1 / 2 / np.sqrt(p1_bottom * (p * dy - p2_bottom)))
                elif shape_bottom == 3:
                    theta_b = atand(-1 / t1_bottom)
                elif shape_bottom == 4:
                    theta_b = atand(0)
                elif shape_bottom == 5:
                    k_B_b_tem = (y_B_b[1] - y_B_b[0]) / (x_B_b[1] - x_B_b[0])
                    theta_b = atand(-1 / k_B_b_tem)
                theta = (atand(-1 / ((y_B_t[i_index + 1] - y_B_t[i_index]) / (x_B_t[i_index + 1] - x_B_t[i_index]))) + theta_b) / 2
            elif y_B_t[i_index + 1] == p * dy and y_B_t[i_index + 1] != 0:
                k1 = (y_B_t[i_index + 1] - y_B_t[i_index]) / (x_B_t[i_index + 1] - x_B_t[i_index])
                k2 = (y_B_t[i_index + 2] - y_B_t[i_index + 1]) / (x_B_t[i_index + 2] - x_B_t[i_index + 1])
                theta = (atand(-1 / k1) + atand(-1 / k2)) / 2
            else:
                k1 = (y_B_t[i_index + 1] - y_B_t[i_index]) / (x_B_t[i_index + 1] - x_B_t[i_index])
                theta = atand(-1 / k1)

        ny = sind(theta)
        nx = cosd(theta)
        zr_1 = 1j / k0 * (nx * nx / Z + Z * ny * ny)
        zr_2 = -1j / k0 * (Z - 1 / Z) * nx * ny
        zr_3 = zr_2
        zr_4 = 1j / k0 * (Z * nx * nx + ny * ny / Z)
        zl_1 = zr_1
        zl_2 = -zr_2
        zl_3 = zl_2
        zl_4 = zr_4
        # 注意: 直接使用Matlab索引 (i_idx+1, x_c±Ix_t(i_idx))
        Z_x1[i_idx + 1, x_c + Ix_t[i_idx]] = zr_1
        Z_x2[i_idx + 1, x_c + Ix_t[i_idx]] = zr_2
        Z_x3[i_idx + 1, x_c + Ix_t[i_idx]] = zr_3
        Z_x4[i_idx + 1, x_c + Ix_t[i_idx]] = zr_4
        Z_x1[i_idx + 1, x_c - Ix_t[i_idx]] = zl_1
        Z_x2[i_idx + 1, x_c - Ix_t[i_idx]] = zl_2
        Z_x3[i_idx + 1, x_c - Ix_t[i_idx]] = zl_3
        Z_x4[i_idx + 1, x_c - Ix_t[i_idx]] = zl_4
        a_x[i_idx + 1, x_c + Ix_t[i_idx] - 1] = Rx_t[i_idx]

        R_x1[i_idx + 1, x_c + Ix_t[i_idx]] = Rx_t[i_idx]
        R_x3[i_idx + 1, x_c + Ix_t[i_idx]] = Rx_t[i_idx] / nx
        R_x2[i_idx + 1, x_c + Ix_t[i_idx]] = ny * Rx_t[i_idx] / nx
        R_x2_c[i_idx + 1, x_c + Ix_t[i_idx]] = 1 - ny * Rx_t[i_idx] / nx
        a_x[i_idx + 1, x_c - Ix_t[i_idx] + 1] = Rx_t[i_idx]

        R_x1[i_idx + 1, x_c - Ix_t[i_idx]] = Rx_t[i_idx]
        R_x3[i_idx + 1, x_c - Ix_t[i_idx]] = Rx_t[i_idx] / nx
        R_x2[i_idx + 1, x_c - Ix_t[i_idx]] = ny * Rx_t[i_idx] / nx
        R_x2_c[i_idx + 1, x_c - Ix_t[i_idx]] = 1 - ny * Rx_t[i_idx] / nx
        # a_x(y, x_l+2 : x_r-2) = 1 -> Python切片右端+1
        x_l = x_c - Ix_t[i_idx]
        x_r = x_c + Ix_t[i_idx]
        # if x_l + 2 <= x_r - 2:
        a_x[i_idx + 1, (x_l + 2) : (x_r - 2) + 1] = 1

        # R_in_d(y, x_l+1:x_r-1) = 1
        # if x_l + 1 <= x_r - 1:
        R_in_d[i_idx + 1, (x_l + 1) : (x_r - 1) + 1] = 1

        i_idx += 1
        
    timeCount('1st for')
    # p2
    i_idx = 0
    for p in range(-1, -nh_bottom, -1):
        if shape_bottom == 1:
            theta = atand((R1_bottom * (p * dy - Hs_bottom)) / (R2_bottom * np.sqrt(R2_bottom ** 2 - (p * dy - Hs_bottom) ** 2)))
        elif shape_bottom == 2:
            theta = atand(-1 / 2 / np.sqrt(p1_bottom * (p * dy - p2_bottom)))
        elif shape_bottom == 3:
            theta = atand(-1 / t1_bottom)
        elif shape_bottom == 4:
            theta = atand(0)
        elif shape_bottom == 5:
            i_index = np.where(y_B_b > p * dy)[0]
            i_index = i_index[-1] if i_index.size > 0 else 0
            if y_B_b[i_index + 1] == p * dy:
                theta = (atand(-1 / ((y_B_b[i_index + 1] - y_B_b[i_index]) / (x_B_b[i_index + 1] - x_B_b[i_index]))) +
                         atand(-1 / ((y_B_b[i_index + 2] - y_B_b[i_index + 1]) / (x_B_b[i_index + 2] - x_B_b[i_index + 1])))) / 2
            else:
                k1 = (y_B_b[i_index + 1] - y_B_b[i_index]) / (x_B_b[i_index + 1] - x_B_b[i_index])
                theta = atand(-1 / k1)

        ny = sind(theta)
        nx = cosd(theta)
        zr_1 = 1j / k0 * (nx * nx / Z + Z * ny * ny)
        zr_2 = -1j / k0 * (Z - 1 / Z) * nx * ny
        zr_3 = zr_2
        zr_4 = 1j / k0 * (Z * nx * nx + ny * ny / Z)
        zl_1 = zr_1
        zl_2 = -zr_2
        zl_3 = zl_2
        zl_4 = zr_4

        Z_x1[y_c + i_idx + 1, x_c + Ix_b[i_idx]] = zr_1
        Z_x2[y_c + i_idx + 1, x_c + Ix_b[i_idx]] = zr_2
        Z_x3[y_c + i_idx + 1, x_c + Ix_b[i_idx]] = zr_3
        Z_x4[y_c + i_idx + 1, x_c + Ix_b[i_idx]] = zr_4
        Z_x1[y_c + i_idx + 1, x_c - Ix_b[i_idx]] = zl_1
        Z_x2[y_c + i_idx + 1, x_c - Ix_b[i_idx]] = zl_2
        Z_x3[y_c + i_idx + 1, x_c - Ix_b[i_idx]] = zl_3
        Z_x4[y_c + i_idx + 1, x_c - Ix_b[i_idx]] = zl_4
        a_x[y_c + i_idx + 1, x_c + Ix_b[i_idx] - 1] = Rx_b[i_idx]
        R_x1[y_c + i_idx + 1, x_c + Ix_b[i_idx]] = Rx_b[i_idx]
        R_x3[y_c + i_idx + 1, x_c + Ix_b[i_idx]] = Rx_b[i_idx] / abs(nx)
        R_x2[y_c + i_idx + 1, x_c + Ix_b[i_idx]] = abs(ny) * Rx_b[i_idx] / abs(nx)
        R_x2_c[y_c + i_idx + 1, x_c + Ix_b[i_idx]] = 1 - abs(ny) * Rx_b[i_idx] / abs(nx)
        a_x[y_c + i_idx + 1, x_c - Ix_b[i_idx] + 1] = Rx_b[i_idx]
        R_x1[y_c + i_idx + 1, x_c - Ix_b[i_idx]] = Rx_b[i_idx]
        R_x3[y_c + i_idx + 1, x_c - Ix_b[i_idx]] = Rx_b[i_idx] / abs(nx)
        R_x2[y_c + i_idx + 1, x_c - Ix_b[i_idx]] = abs(ny) * Rx_b[i_idx] / abs(nx)
        R_x2_c[y_c + i_idx + 1, x_c - Ix_b[i_idx]] = 1 - abs(ny) * Rx_b[i_idx] / abs(nx)
        x_l = x_c - Ix_b[i_idx]
        x_r = x_c + Ix_b[i_idx]
        a_x[y_c + i_idx + 1, (x_l + 2) : (x_r - 2) + 1] = 1
        R_in_d[y_c + i_idx + 1, (x_l + 1) : (x_r - 1) + 1] = 1
        i_idx += 1

    timeCount('2 for end')
    # p3
    a_y = np.full((Ny + 1, Nx + 1), np.inf, dtype=float)
    R_y1 = np.zeros((Ny + 1, Nx + 1), dtype=float)
    R_y2 = np.zeros((Ny + 1, Nx + 1), dtype=float)
    R_y3 = np.zeros((Ny + 1, Nx + 1), dtype=complex)
    R_y1_c = np.zeros((Ny + 1, Nx + 1), dtype=float)
    R_y1_n = np.zeros((Ny + 1, Nx + 1), dtype=float)
    Z_y1 = np.zeros((Ny + 1, Nx + 1), dtype=complex)
    Z_y2 = np.zeros((Ny + 1, Nx + 1), dtype=complex)
    Z_y3 = np.zeros((Ny + 1, Nx + 1), dtype=complex)
    Z_y4 = np.zeros((Ny + 1, Nx + 1), dtype=complex)

    if shape_top in (1, 2, 3):
        xxx_t_edge = X_t
    elif shape_top == 4:
        xxx_t_edge = W_half
    else:
        xxx_t_edge = 0.0
    from func_loc_res_coef import func_loc_res_coef as _flrc
    Ix_t_edge, _ = _flrc(v_roundn(xxx_t_edge / dx))

    # 从第401行到第427行，对应matlab的第401行到第427行 (构造yyy_t)
    xxx_t = np.arange(0, nw_half) * dx
    if shape_top == 1:
        yyy_t = R2_top * np.sqrt(1 - xxx_t ** 2 / R1_top ** 2) - Hs_top
    elif shape_top == 2:
        yyy_t = p1_top * xxx_t ** 2 + p2_top
    elif shape_top == 3:
        yyy_t = t1_top * xxx_t + t2_top
    elif shape_top == 4:
        yyy_t = H_top * np.ones_like(xxx_t)
    elif shape_top == 5:
        k_B_t = np.zeros(n_seg_t)
        c_B_t = np.zeros(n_seg_t)
        yyy_t_list = []
        i = 0
        k_B_t[i] = (y_B_t[i + 1] - y_B_t[i]) / (x_B_t[i + 1] - x_B_t[i])
        c_B_t[i] = y_B_t[i] - k_B_t[i] * x_B_t[i]
        xxx_t_tentative = np.arange(0, int(np.floor(v_roundn(x_B_t[i + 1] / dx))) + 1) * dx
        yyy_t_tentative = k_B_t[i] * xxx_t_tentative + c_B_t[i]
        yyy_t_list.append(yyy_t_tentative)
        for i in range(1, n_seg_t):
            k_B_t[i] = (y_B_t[i + 1] - y_B_t[i]) / (x_B_t[i + 1] - x_B_t[i])
            c_B_t[i] = y_B_t[i] - k_B_t[i] * x_B_t[i]
            xa = int(np.floor(v_roundn(x_B_t[i] / dx)))
            xb = int(np.floor(v_roundn(x_B_t[i + 1] / dx)))
            if xb >= xa + 1:
                xxx_t_tentative = np.arange(xa + 1, xb + 1) * dx
                yyy_t_tentative = k_B_t[i] * xxx_t_tentative + c_B_t[i]
            else:
                yyy_t_tentative = np.array([0.0])
            yyy_t_list.append(yyy_t_tentative)
        yyy_t = np.concatenate(yyy_t_list)[:-1]
    # yyy_t(1:Ix_t_edge) = H_top
    if Ix_t_edge > 0:
        yyy_t[: Ix_t_edge] = H_top

    # 从第429行到第469行，对应matlab的第429行到第469行 (构造yyy_b与Ix_b_edge)
    xxx_b = np.arange(0, nw_half) * dx
    if shape_bottom in (1, 2, 3):
        xxx_b_edge = X_b
    elif shape_bottom == 4:
        xxx_b_edge = W_half
    else:
        xxx_b_edge = 0.0
    Ix_b_edge, _ = _flrc(v_roundn(xxx_b_edge / dx))

    if shape_bottom == 1:
        yyy_b = -R2_bottom * np.sqrt(1 - xxx_b ** 2 / R1_bottom ** 2) + Hs_bottom
    elif shape_bottom == 2:
        yyy_b = p1_bottom * xxx_b ** 2 + p2_bottom
    elif shape_bottom == 3:
        yyy_b = t1_bottom * xxx_b + t2_bottom
    elif shape_bottom == 4:
        yyy_b = -H_bottom * np.ones_like(xxx_b)
    elif shape_bottom == 5:
        k_B_b = np.zeros(n_seg_b)
        c_B_b = np.zeros(n_seg_b)
        yyy_b_list = []
        i = n_seg_b - 1
        k_B_b[i] = (y_B_b[i + 1] - y_B_b[i]) / (x_B_b[i + 1] - x_B_b[i])
        c_B_b[i] = y_B_b[i] - k_B_b[i] * x_B_b[i]
        xxx_b_tentative = np.arange(0, int(np.floor(v_roundn(x_B_b[i] / dx))) + 1) * dx
        yyy_b_tentative = k_B_b[i] * xxx_b_tentative + c_B_b[i]
        yyy_b_list.append(yyy_b_tentative)
        for i in range(n_seg_b - 2, -1, -1):
            k_B_b[i] = (y_B_b[i + 1] - y_B_b[i]) / (x_B_b[i + 1] - x_B_b[i])
            c_B_b[i] = y_B_b[i] - k_B_b[i] * x_B_b[i]
            xa = int(np.floor(v_roundn(x_B_b[i + 1] / dx)))
            xb = int(np.floor(v_roundn(x_B_b[i] / dx)))
            if xb >= xa + 1:
                xxx_b_tentative = np.arange(xa + 1, xb + 1) * dx
                yyy_b_tentative = k_B_b[i] * xxx_b_tentative + c_B_b[i]
            else:
                yyy_b_tentative = np.array([0.0])
            yyy_b_list.append(yyy_b_tentative)
        yyy_b = np.concatenate(yyy_b_list)[:-1]
    if Ix_b_edge > 0:
        yyy_b[: Ix_b_edge] = -H_bottom

    Iy_t, Ry_t = _flrc(v_roundn(yyy_t / dy))
    Iy_b, Ry_b = _flrc(v_roundn(yyy_b / dy))


    # 从第472行到第535行，对应matlab的第472行到第535行 (平台对偶阻抗与边界系数-左半/右半)
    j = 0
    for q in range(0, Ix_b_edge):
        timeCount(f'in 3 for, j={j}')
        # 顶部theta
        if shape_top == 1:
            theta = atand(R1_top * np.sqrt(R1_top ** 2 - (q * dx) ** 2) / (R2_top * q * dx))
        elif shape_top == 2:
            theta = atand(-1 / 2 / p1_top / (q * dx))
        elif shape_top == 3:
            if q * dx < X_b:
                theta = atand(np.inf)
            elif q * dx == X_b:
                theta = (atand(np.inf) + atand(-1 / t1_top)) / 2
            else:
                theta = atand(-1 / t1_top)
        else:
            theta = atand(np.inf)
        ny = sind(theta)
        nx = cosd(theta)
        zr_1 = 1j / k0 * (nx * nx / Z + Z * ny * ny)
        zr_2 = -1j / k0 * (Z - 1 / Z) * nx * ny
        zr_3 = zr_2
        zr_4 = 1j / k0 * (Z * nx * nx + ny * ny / Z)
        zl_1 = zr_1
        zl_2 = -zr_2
        zl_3 = zl_2
        zl_4 = zr_4
        # y_c - Iy_t(j+1), x_c ± j
        y_up = y_c - Iy_t[j]
        x_r = x_c + j
        x_l = x_c - j
        Z_y1[y_up, x_r] = zr_1
        Z_y2[y_up, x_r] = zr_2
        Z_y3[y_up, x_r] = zr_3
        Z_y4[y_up, x_r] = zr_4
        Z_y1[y_up, x_l] = zl_1
        Z_y2[y_up, x_l] = zl_2
        Z_y3[y_up, x_l] = zl_3
        Z_y4[y_up, x_l] = zl_4
        a_y[y_up + 1, x_r] = Ry_t[j]
        R_y2[y_up, x_r] = Ry_t[j]
        R_y3[y_up, x_r] = Ry_t[j] / ny
        R_y1[y_up, x_r] = nx * Ry_t[j] / ny
        R_y1_c[y_up, x_r] = 1 - nx * Ry_t[j] / ny
        a_y[y_up + 1, x_l] = Ry_t[j]
        R_y2[y_up, x_l] = Ry_t[j]
        R_y3[y_up, x_l] = Ry_t[j] / ny
        R_y1[y_up, x_l] = nx * Ry_t[j] / ny
        R_y1_c[y_up, x_l] = 1 - nx * Ry_t[j] / ny
        # 底部theta = atand(-Inf)
        theta = atand(-np.inf)
        ny = sind(theta)
        nx = cosd(theta)
        zr_1 = 1j / k0 * (nx * nx / Z + Z * ny * ny)
        zr_2 = -1j / k0 * (Z - 1 / Z) * nx * ny
        zr_3 = zr_2
        zr_4 = 1j / k0 * (Z * nx * nx + ny * ny / Z)
        zl_1 = zr_1
        zl_2 = -zr_2
        zl_3 = zl_2
        zl_4 = zr_4
        y_dn = y_c + Iy_b[j]
        Z_y1[y_dn, x_r] = zr_1
        Z_y2[y_dn, x_r] = zr_2
        Z_y3[y_dn, x_r] = zr_3
        Z_y4[y_dn, x_r] = zr_4
        Z_y1[y_dn, x_l] = zl_1
        Z_y2[y_dn, x_l] = zl_2
        Z_y3[y_dn, x_l] = zl_3
        Z_y4[y_dn, x_l] = zl_4
        a_y[y_dn - 1, x_r] = Ry_b[j]
        R_y2[y_dn, x_r] = Ry_b[j]
        R_y3[y_dn, x_r] = Ry_b[j] / abs(ny)
        R_y1[y_dn, x_r] = abs(nx) * Ry_b[j] / abs(ny)
        R_y1_c[y_dn, x_r] = 1 - abs(nx) * Ry_b[j] / abs(ny)
        a_y[y_dn - 1, x_l] = Ry_b[j]
        R_y2[y_dn, x_l] = Ry_b[j]
        R_y3[y_dn, x_l] = Ry_b[j] / abs(ny)
        R_y1[y_dn, x_l] = abs(nx) * Ry_b[j] / abs(ny)
        R_y1_c[y_dn, x_l] = 1 - abs(nx) * Ry_b[j] / abs(ny)
        # a_y(y_up+2:y_dn-2, x_±j) = 1
        # if (y_up + 2) <= (y_dn - 2):
        a_y[(y_up + 2) : (y_dn - 2) + 1, x_l] = 1
        a_y[(y_up + 2) : (y_dn - 2) + 1, x_r] = 1
        j += 1

    # 从第537行到第592行，对应matlab的第537行到第592行 (右半区 j=Ix_b_edge:...)
    j = Ix_b_edge
    for q in range(Ix_b_edge, nw_half):
        timeCount(f'in 4 for, q={q}')
        if shape_top == 1:
            theta = atand(R1_top * np.sqrt(R1_top ** 2 - (q * dx) ** 2) / (R2_top * q * dx))
        elif shape_top == 2:
            theta = atand(-1 / 2 / p1_top / (q * dx))
        elif shape_top == 3:
            theta = atand(-1 / t1_top)
        elif shape_top == 4:
            theta = atand(np.inf)
        elif shape_top == 5:
            i_index = np.where(x_B_t >= q * dx)[0]
            i_index = i_index[0] if i_index.size > 0 else 0
            if x_B_t[i_index] == q * dx and q * dx == 0:
                theta = atand(np.inf)
            elif x_B_t[i_index] == q * dx and q * dx != 0:
                kcur = (y_B_t[i_index + 1] - y_B_t[i_index]) / (x_B_t[i_index + 1] - x_B_t[i_index])
                theta = (90 + atand(-1 / kcur)) / 2
            else:
                kprev = (y_B_t[i_index] - y_B_t[i_index - 1]) / (x_B_t[i_index] - x_B_t[i_index - 1])
                theta = atand(-1 / kprev)
                if theta == -90:
                    theta = 90
        ny = sind(theta)
        nx = cosd(theta)
        zr_1 = 1j / k0 * (nx * nx / Z + Z * ny * ny)
        zr_2 = -1j / k0 * (Z - 1 / Z) * nx * ny
        zr_3 = zr_2
        zr_4 = 1j / k0 * (Z * nx * nx + ny * ny / Z)
        zl_1 = zr_1
        zl_2 = -zr_2
        zl_3 = zl_2
        zl_4 = zr_4
        y_up = y_c - Iy_t[j + 1 - 1]
        x_r = x_c + j
        x_l = x_c - j
        Z_y1[y_up, x_r] = zr_1
        Z_y2[y_up, x_r] = zr_2
        Z_y3[y_up, x_r] = zr_3
        Z_y4[y_up, x_r] = zr_4
        Z_y1[y_up, x_l] = zl_1
        Z_y2[y_up, x_l] = zl_2
        Z_y3[y_up, x_l] = zl_3
        Z_y4[y_up, x_l] = zl_4
        a_y[y_up + 1, x_r] = Ry_t[j + 1 - 1]
        R_y2[y_up, x_r] = Ry_t[j + 1 - 1]
        R_y3[y_up, x_r] = Ry_t[j + 1 - 1] / ny
        R_y1[y_up, x_r] = nx * Ry_t[j + 1 - 1] / ny
        R_y1_c[y_up, x_r] = 1 - nx * Ry_t[j + 1 - 1] / ny
        a_y[y_up + 1, x_l] = Ry_t[j + 1 - 1]
        R_y2[y_up, x_l] = Ry_t[j + 1 - 1]
        R_y3[y_up, x_l] = Ry_t[j + 1 - 1] / ny
        R_y1[y_up, x_l] = nx * Ry_t[j + 1 - 1] / ny
        R_y1_c[y_up, x_l] = 1 - nx * Ry_t[j + 1 - 1] / ny

        # bottom
        if shape_bottom == 1:
            theta = atand(-R1_bottom * np.sqrt(R1_bottom ** 2 - (q * dx) ** 2) / (R2_bottom * q * dx))
        elif shape_bottom == 2:
            theta = atand(-1 / 2 / p1_bottom / (q * dx))
        elif shape_bottom == 3:
            theta = atand(-1 / t1_bottom)
        elif shape_bottom == 4:
            theta = atand(-np.inf)
        elif shape_bottom == 5:
            x_B_b_tem = x_B_b[::-1]
            k_B_b_tem = np.diff(y_B_b[::-1]) / np.diff(x_B_b[::-1])
            idxs = np.where(x_B_b_tem >= q * dx)[0]
            i_index = idxs[0] if idxs.size > 0 else 0
            if x_B_b_tem[i_index] == q * dx and q * dx == 0:
                theta = atand(-np.inf)
            elif x_B_b_tem[i_index] == q * dx and q * dx != 0:
                theta = (-90 + atand(-1 / k_B_b_tem[i_index])) / 2
            else:
                theta = atand(-1 / k_B_b_tem[i_index - 1])
                if theta == 90:
                    theta = -90
        ny = sind(theta)
        nx = cosd(theta)
        zr_1 = 1j / k0 * (nx * nx / Z + Z * ny * ny)
        zr_2 = -1j / k0 * (Z - 1 / Z) * nx * ny
        zr_3 = zr_2
        zr_4 = 1j / k0 * (Z * nx * nx + ny * ny / Z)
        zl_1 = zr_1
        zl_2 = -zr_2
        zl_3 = zl_2
        zl_4 = zr_4
        y_dn = y_c + Iy_b[j + 1 - 1]
        Z_y1[y_dn, x_r] = zr_1
        Z_y2[y_dn, x_r] = zr_2
        Z_y3[y_dn, x_r] = zr_3
        Z_y4[y_dn, x_r] = zr_4
        Z_y1[y_dn, x_l] = zl_1
        Z_y2[y_dn, x_l] = zl_2
        Z_y3[y_dn, x_l] = zl_3
        Z_y4[y_dn, x_l] = zl_4
        a_y[y_dn - 1, x_r] = Ry_b[j + 1 - 1]
        R_y2[y_dn, x_r] = Ry_b[j + 1 - 1]
        R_y3[y_dn, x_r] = Ry_b[j + 1 - 1] / abs(ny)
        R_y1[y_dn, x_r] = abs(nx) * Ry_b[j + 1 - 1] / abs(ny)
        R_y1_c[y_dn, x_r] = 1 - abs(nx) * Ry_b[j + 1 - 1] / abs(ny)
        a_y[y_dn - 1, x_l] = Ry_b[j + 1 - 1]
        R_y2[y_dn, x_l] = Ry_b[j + 1 - 1]
        R_y3[y_dn, x_l] = Ry_b[j + 1 - 1] / abs(ny)
        R_y1[y_dn, x_l] = abs(nx) * Ry_b[j + 1 - 1] / abs(ny)
        R_y1_c[y_dn, x_l] = 1 - abs(nx) * Ry_b[j + 1 - 1] / abs(ny)
        if (y_up + 2) <= (y_dn - 2):
            a_y[(y_up + 2) : (y_dn - 2) + 1, x_l] = 1
            a_y[(y_up + 2) : (y_dn - 2) + 1, x_r] = 1
        j += 1

    # 从第594行到第606行，对应matlab的第594行到第606行 (R_x2>1 和 R_y1>1 的修正)
    rx_pos = np.argwhere(R_x2 > 1)
    for (ix, iy) in rx_pos:
        R_x2_c[ix, iy] = 0
        R_x2_n[ix, iy] = R_x1[ix, iy] * (R_x2[ix, iy] - 1) / R_x2[ix, iy]
        R_x3[ix, iy] = R_x3[ix, iy] * (1 - (R_x2[ix, iy] - 1) / R_x2[ix, iy])
        R_x2[ix, iy] = 1 - R_x2_n[ix, iy]

    ry_pos = np.argwhere(R_y1 > 1)
    for (ix, iy) in ry_pos:
        R_y1_c[ix, iy] = 0
        R_y1_n[ix, iy] = R_y2[ix, iy] * (R_y1[ix, iy] - 1) / R_y1[ix, iy]
        R_y3[ix, iy] = R_y3[ix, iy] * (1 - (R_y1[ix, iy] - 1) / R_y1[ix, iy])
        R_y1[ix, iy] = 1 - R_y1_n[ix, iy]

    R_B_d1 = np.ceil(R_x1)
    R_B_d2 = np.ceil(R_y2)

    # 从第609行到第636行，对应matlab的第609行到第636行 (系数与数组预分配)
    rx = ds / dx / dx
    ry = ds / dy / dy
    cx = rx / (4 * 1j * k0)
    cy = ry / (4 * 1j * k0)

    # 下面大量向量尺寸按Matlab原样分配 (均为1基索引区域将使用1..Nx+1或1..Ny+1)
    bc_t_A1 = np.zeros(2 * Nx + 1, dtype=np.complex128)
    bc_b_A1 = np.zeros(2 * Nx + 1, dtype=np.complex128)
    d_t_xy_A1 = np.zeros(Nx)
    d_b_xy_A1 = np.zeros(Nx + 2)
    d_t_yx_A1 = np.zeros(Nx + 2)
    d_b_yx_A1 = np.zeros(Nx)
    bc_t_A1_u = np.zeros(2 * Nx + 1)
    bc_b_A1_u = np.zeros(2 * Nx + 1)
    d_t_xy_A1_u = np.zeros(Nx)
    d_b_xy_A1_u = np.zeros(Nx + 2)
    d_t_yx_A1_u = np.zeros(Nx + 2)
    d_b_yx_A1_u = np.zeros(Nx)
    bc_t_A1_d = np.zeros(2 * Nx + 1)
    bc_b_A1_d = np.zeros(2 * Nx + 1)
    d_t_xy_A1_d = np.zeros(Nx)
    d_b_xy_A1_d = np.zeros(Nx + 2)
    d_t_yx_A1_d = np.zeros(Nx + 2)
    d_b_yx_A1_d = np.zeros(Nx)
    bc_t_A2 = np.zeros((2 * Ny + 1, 1))
    bc_b_A2 = np.zeros((2 * Ny + 1, 1))
    d_t_xy_A2 = np.zeros((Ny, 1))
    d_b_xy_A2 = np.zeros((Ny + 2, 1))
    d_t_yx_A2 = np.zeros((Ny + 2, 1))
    d_b_yx_A2 = np.zeros((Ny, 1))
    bc_t_A2_l = np.zeros((2 * Ny + 1, 1))
    bc_b_A2_l = np.zeros((2 * Ny + 1, 1))
    d_t_xy_A2_l = np.zeros((Ny, 1))
    d_b_xy_A2_l = np.zeros((Ny + 2, 1))
    d_t_yx_A2_l = np.zeros((Ny + 2, 1))
    d_b_yx_A2_l = np.zeros((Ny, 1))
    bc_t_A2_r = np.zeros((2 * Ny + 1, 1))
    bc_b_A2_r = np.zeros((2 * Ny + 1, 1))
    d_t_xy_A2_r = np.zeros((Ny, 1))
    d_b_xy_A2_r = np.zeros((Ny + 2, 1))
    d_t_yx_A2_r = np.zeros((Ny + 2, 1))
    d_b_yx_A2_r = np.zeros((Ny, 1))
    bx_1 = np.zeros((Nx + 1, 1), dtype=complex)
    bx_1_u = np.zeros((Nx + 1, 1), dtype=complex)
    bx_1_d = np.zeros((Nx + 1, 1), dtype=complex)
    by_1 = np.zeros_like(bx_1)
    by_1_u = np.zeros_like(bx_1_u)
    by_1_d = np.zeros_like(bx_1_d)
    bx_2 = np.zeros((Ny + 1, 1), dtype=complex)
    bx_2_l = np.zeros((Ny + 1, 1), dtype=complex)
    bx_2_r = np.zeros((Ny + 1, 1), dtype=complex)
    by_2 = np.zeros_like(bx_2)
    by_2_l = np.zeros_like(bx_2_l)
    by_2_r = np.zeros_like(bx_2_r)

    # 从第638行到第660行，对应matlab的第638行到第660行 (C系数计算)
    C_d_x = (dx * R_x3 - Z_x1) * (dx * R_x3 - Z_x4) - Z_x2 * Z_x3
    C1_x = -(Z_x1 * (dx * R_x3 - Z_x4) + Z_x2 * Z_x3) / C_d_x
    C3_x = -(Z_x4 * (dx * R_x3 - Z_x1) + Z_x2 * Z_x3) / C_d_x
    C2_x = -(Z_x3 * (dx * R_x3 - Z_x4) + Z_x3 * Z_x4) / C_d_x
    C4_x = -(Z_x2 * (dx * R_x3 - Z_x1) + Z_x1 * Z_x2) / C_d_x
    C_d_y = (dy * R_y3 - Z_y1) * (dy * R_y3 - Z_y4) - Z_y2 * Z_y3
    C1_y = -(Z_y1 * (dy * R_y3 - Z_y4) + Z_y2 * Z_y3) / C_d_y
    C3_y = -(Z_y4 * (dy * R_y3 - Z_y1) + Z_y2 * Z_y3) / C_d_y
    C2_y = -(Z_y3 * (dy * R_y3 - Z_y4) + Z_y3 * Z_y4) / C_d_y
    C4_y = -(Z_y2 * (dy * R_y3 - Z_y1) + Z_y1 * Z_y2) / C_d_y
    C1_x *= R_B_d1
    C2_x *= R_B_d1
    C3_x *= R_B_d1
    C4_x *= R_B_d1
    C1_y *= R_B_d2
    C2_y *= R_B_d2
    C3_y *= R_B_d2
    C4_y *= R_B_d2
    for M in (C1_x, C2_x, C3_x, C4_x, C1_y, C2_y, C3_y, C4_y):
        M[np.isnan(M)] = 0

    # 从第663行到第679行，对应matlab的第663行到第679行 (边界条件PEC/PMC调整)
    if BC == 2:  # PEC
        C1_x = 0 * R_B_d1
        C2_x = 0 * R_B_d1
        C3_x = 0 * R_B_d1
        C4_x = 0 * R_B_d1
        C1_y = 0 * R_B_d2
        C2_y = 0 * R_B_d2
        C3_y = 0 * R_B_d2
        C4_y = 0 * R_B_d2
    elif BC == 3:  # PMC
        C1_x = R_B_d1
        C2_x = 0 * R_B_d1
        C3_x = R_B_d1
        C4_x = 0 * R_B_d1
        C1_y = R_B_d2
        C2_y = 0 * R_B_d2
        C3_y = R_B_d2
        C4_y = 0 * R_B_d2

    # 从第682行到第739行，对应matlab的第682行到第739行 (初始平面与源设置，以及主推进循环)
    Tx_vertical = TX[1] - H_bottom
    Tx_horizontal = TX[0]
    w0 = k_w0 * lambda_
    u_g_c = np.zeros((Ny + 1, Nx + 1), dtype=complex)
    C = 1
    n_GB_x = 1
    n_GB_y = n_GB_x
    n_GB = n_GB_x * n_GB_y
    w0_x = np.zeros(n_GB)
    w0_y = np.zeros(n_GB)
    z_far = 0
    for ii in range(n_GB_x):
        w0_x[(ii) * n_GB_x : (ii + 1) * n_GB_x] = w0
        w0_y[(ii) * n_GB_y : (ii + 1) * n_GB_y] = w0
    # 下面的w_x等仅用于形成常量，维度按标量处理
    w_x = w0_x * np.sqrt(1 + (lambda_ * z_far / np.pi) ** 2 / w0_x ** 4)
    R_x = z_far + np.pi**2 / lambda_**2 / z_far * w0_x ** 4 if z_far != 0 else np.inf
    phi0_x = np.arctan(lambda_ * z_far / np.pi / w0_x ** 2) if z_far != 0 else 0.0
    w_y = w0_y * np.sqrt(1 + (lambda_ * z_far / np.pi) ** 2 / w0_y ** 4)
    R_y = z_far + np.pi**2 / lambda_**2 / z_far * w0_y ** 4 if z_far != 0 else np.inf
    phi0_y = np.arctan(lambda_ * z_far / np.pi / w0_y ** 2) if z_far != 0 else 0.0
    iii = 0
    for ny_row in range(nh_top, -nh_bottom-1, -1):
        jjj = 0
        for nx_col in range(-nw_half, nw_half + 1):
            u_g_c[iii, jjj] = (np.sqrt(2.0 / np.pi / (w_x * w_y)) * np.exp( -( (nx_col * dx - Tx_horizontal) ** 2 / w_x ** 2 + (ny_row * dy - Tx_vertical) ** 2 / w_y ** 2) - 1j * k0 * z_far - 1j * np.pi * ( (nx_col * dx - Tx_horizontal) ** 2 / lambda_ / R_x +
                            (ny_row * dy - Tx_vertical) ** 2 / lambda_ / R_y) + 1j * phi0_x / 2 + 1j * phi0_y / 2) * C)
            jjj += 1
        iii += 1
        
    if (u_g_c.imag == 0).all():
        u_g_c = u_g_c.real

    Source_x = 0
    Source_y = u_g_c / np.max(np.abs(u_g_c))
    Ex = Source_x * R_in_d
    Ey = Source_y * R_in_d

    # 初始边界点修正
    y_B_c = np.where(R_B_d2[:, x_c] == 1)[0]
    # 注意: y_B_c含0索引项，Matlab中使用(y_B_c(1), x_c)与(y_B_c(2), x_c)
    Ex[y_B_c[0], x_c] = C1_y[y_B_c[0], x_c] * Ex[y_B_c[0] + 1, x_c] + C2_y[y_B_c[0], x_c] * Ey[y_B_c[0] + 1, x_c]
    Ey[y_B_c[0], x_c] = C3_y[y_B_c[0], x_c] * Ey[y_B_c[0] + 1, x_c] + C4_y[y_B_c[0], x_c] * Ex[y_B_c[0] + 1, x_c]
    Ex[y_B_c[1], x_c] = C1_y[y_B_c[1], x_c] * Ex[y_B_c[1] - 1, x_c] + C2_y[y_B_c[1], x_c] * Ey[y_B_c[1] - 1, x_c]
    Ey[y_B_c[1], x_c] = C3_y[y_B_c[1], x_c] * Ey[y_B_c[1] - 1, x_c] + C4_y[y_B_c[1], x_c] * Ex[y_B_c[1] - 1, x_c]



    timeCount(f'in 5 for')
    # 从第889行到第912行这部分python代码，对应的是matlab的第889行到第912行
    for jj in range(1, nw_half):
        x_r = x_c + jj
        x_l = x_c - jj
        y_B_r = np.where(R_B_d2[:, x_r] == 1)[0]
        y_B_l = np.where(R_B_d2[:, x_l] == 1)[0]
        Ex[y_B_r[0], x_r] = ( C1_y[y_B_r[0], x_r] * ( R_y1[y_B_r[0], x_r] * Ex[y_B_r[0] + 1, x_r - 1] + R_y1_c[y_B_r[0], x_r] * Ex[y_B_r[0] + 1, x_r] + R_y1_n[y_B_r[0], x_r] * Ex[y_B_r[0], x_r - 1]) + C2_y[y_B_r[0], x_r] * ( R_y1[y_B_r[0], x_r] * Ey[y_B_r[0] + 1, x_r - 1] + R_y1_c[y_B_r[0], x_r] * Ey[y_B_r[0] + 1, x_r] + R_y1_n[y_B_r[0], x_r] * Ey[y_B_r[0], x_r - 1]))
        Ey[y_B_r[0], x_r] = ( C3_y[y_B_r[0], x_r] * ( R_y1[y_B_r[0], x_r] * Ey[y_B_r[0] + 1, x_r - 1] + R_y1_c[y_B_r[0], x_r] * Ey[y_B_r[0] + 1, x_r] + R_y1_n[y_B_r[0], x_r] * Ey[y_B_r[0], x_r - 1]) + C4_y[y_B_r[0], x_r] * ( R_y1[y_B_r[0], x_r] * Ex[y_B_r[0] + 1, x_r - 1] + R_y1_c[y_B_r[0], x_r] * Ex[y_B_r[0] + 1, x_r] + R_y1_n[y_B_r[0], x_r] * Ex[y_B_r[0], x_r - 1]))
        Ex[y_B_r[1], x_r] = ( C1_y[y_B_r[1], x_r] * ( R_y1[y_B_r[1], x_r] * Ex[y_B_r[1] - 1, x_r - 1] + R_y1_c[y_B_r[1], x_r] * Ex[y_B_r[1] - 1, x_r] + R_y1_n[y_B_r[1], x_r] * Ex[y_B_r[1], x_r - 1]) + C2_y[y_B_r[1], x_r] * ( R_y1[y_B_r[1], x_r] * Ey[y_B_r[1] - 1, x_r - 1] + R_y1_c[y_B_r[1], x_r] * Ey[y_B_r[1] - 1, x_r] + R_y1_n[y_B_r[1], x_r] * Ey[y_B_r[1], x_r - 1]))
        Ey[y_B_r[1], x_r] = ( C3_y[y_B_r[1], x_r] * ( R_y1[y_B_r[1], x_r] * Ey[y_B_r[1] - 1, x_r - 1] + R_y1_c[y_B_r[1], x_r] * Ey[y_B_r[1] - 1, x_r] + R_y1_n[y_B_r[1], x_r] * Ey[y_B_r[1], x_r - 1]) + C4_y[y_B_r[1], x_r] * ( R_y1[y_B_r[1], x_r] * Ex[y_B_r[1] - 1, x_r - 1] + R_y1_c[y_B_r[1], x_r] * Ex[y_B_r[1] - 1, x_r] + R_y1_n[y_B_r[1], x_r] * Ex[y_B_r[1], x_r - 1]))
        Ex[y_B_l[0], x_l] = ( C1_y[y_B_l[0], x_l] * ( R_y1[y_B_l[0], x_l] * Ex[y_B_l[0] + 1, x_l + 1] + R_y1_c[y_B_l[0], x_l] * Ex[y_B_l[0] + 1, x_l] + R_y1_n[y_B_l[0], x_l] * Ex[y_B_l[0], x_l + 1]) + C2_y[y_B_l[0], x_l] * ( R_y1[y_B_l[0], x_l] * Ey[y_B_l[0] + 1, x_l + 1] + R_y1_c[y_B_l[0], x_l] * Ey[y_B_l[0] + 1, x_l] + R_y1_n[y_B_l[0], x_l] * Ey[y_B_l[0], x_l + 1]))
        Ey[y_B_l[0], x_l] = ( C3_y[y_B_l[0], x_l] * ( R_y1[y_B_l[0], x_l] * Ey[y_B_l[0] + 1, x_l + 1] + R_y1_c[y_B_l[0], x_l] * Ey[y_B_l[0] + 1, x_l] + R_y1_n[y_B_l[0], x_l] * Ey[y_B_l[0], x_l + 1]) + C4_y[y_B_l[0], x_l] * ( R_y1[y_B_l[0], x_l] * Ex[y_B_l[0] + 1, x_l + 1] + R_y1_c[y_B_l[0], x_l] * Ex[y_B_l[0] + 1, x_l] + R_y1_n[y_B_l[0], x_l] * Ex[y_B_l[0], x_l + 1]))
        Ex[y_B_l[1], x_l] = ( C1_y[y_B_l[1], x_l] * ( R_y1[y_B_l[1], x_l] * Ex[y_B_l[1] - 1, x_l + 1] + R_y1_c[y_B_l[1], x_l] * Ex[y_B_l[1] - 1, x_l] + R_y1_n[y_B_l[1], x_l] * Ex[y_B_l[1], x_l + 1]) + C2_y[y_B_l[1], x_l] * ( R_y1[y_B_l[1], x_l] * Ey[y_B_l[1] - 1, x_l + 1] + R_y1_c[y_B_l[1], x_l] * Ey[y_B_l[1] - 1, x_l] + R_y1_n[y_B_l[1], x_l] * Ey[y_B_l[1], x_l + 1]))
        Ey[y_B_l[1], x_l] = ( C3_y[y_B_l[1], x_l] * ( R_y1[y_B_l[1], x_l] * Ey[y_B_l[1] - 1, x_l + 1] + R_y1_c[y_B_l[1], x_l] * Ey[y_B_l[1] - 1, x_l] + R_y1_n[y_B_l[1], x_l] * Ey[y_B_l[1], x_l + 1]) + C4_y[y_B_l[1], x_l] * ( R_y1[y_B_l[1], x_l] * Ex[y_B_l[1] - 1, x_l + 1] + R_y1_c[y_B_l[1], x_l] * Ex[y_B_l[1] - 1, x_l] + R_y1_n[y_B_l[1], x_l] * Ex[y_B_l[1], x_l + 1]))


    # 辅助: 构造带状对角矩阵 (2*(Nx+1) 或 2*(Ny+1) 维)
    def build_block_diag_A(d_c_main, d_t_off, d_b_off, bc_t, bc_b, d_t_xy, d_b_xy, d_b_yx, d_t_yx, n_block):
        # n_block = Nx 或 Ny
        size = 2 * (n_block + 1)
        A = np.zeros((size, size), dtype=complex)
        # 主对角: [d_c d_c]
        main_vec = np.concatenate([d_c_main, d_c_main]).astype(complex)
        np.fill_diagonal(A, main_vec)
        # 上一对角(+1): ([d_t 0 d_t] - bc_t)
        up_vec = np.concatenate([d_t_off, [0], d_t_off]).astype(complex) - bc_t.astype(complex)
        for k in range(size - 1):
            A[k, k + 1] += up_vec[k]
        # 下一对角(-1): ([d_b 0 d_b] - bc_b)
        lo_vec = np.concatenate([d_b_off, [0], d_b_off]).astype(complex) - bc_b.astype(complex)
        for k in range(1, size):
            A[k, k - 1] += lo_vec[k - 1]
        # 跨块对角(+n_block+2): d_t_xy
        k_off = n_block + 2
        for k in range(size - k_off):
            A[k, k + k_off] += complex(d_t_xy[k])
        # 跨块对角(+n_block): d_b_xy
        k_off2 = n_block
        for k in range(size - k_off2):
            A[k, k + k_off2] += complex(d_b_xy[k])
        # 跨块对角(-n_block-2): d_b_yx
        for k in range(k_off):
            # positions: (k + k_off, k)
            if k + k_off < size:
                A[k + k_off, k] += complex(d_b_yx[k])
        # 跨块对角(-n_block): d_t_yx
        for k in range(k_off2):
            if k + k_off2 < size:
                A[k + k_off2, k] += complex(d_t_yx[k])
        return A

    # 分配工作矩阵与结果容器
    Wx_m = np.zeros((Ny + 1, Nx + 1), dtype=complex)
    Wy_m = np.zeros((Ny + 1, Nx + 1), dtype=complex)
    Ex_all = np.zeros((Ny + 1, Nx + 1, Nz + 1), dtype=complex)
    Ey_all = np.zeros((Ny + 1, Nx + 1, Nz + 1), dtype=complex)
    Ex_all[:, :, 0] = Ex
    Ey_all[:, :, 0] = Ey
    Wx = Ex * np.exp(1j * k0 * s_start)
    Wy = Ey * np.exp(1j * k0 * s_start)

    # 从第471行到第739行这部分python代码，对应的是matlab的第471行到第739行
    # 主推进循环
    for z in range(1, Nz + 1):
        timeCount(f'in z for, j={j}')
        # ---------- A1 中央行 ---------- !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!! iindices here trig bug
        d_c_A1 = 1 + (2 * cx) / a_x[y_c, :]
        d_t_A1 = (-2 * cx) / (1 + a_x[y_c, 0 : Nx]) / np.concatenate([np.ones(x_c), a_x[y_c, x_c + 1 : Nx + 1]])
        d_b_A1 = (-2 * cx) / (1 + a_x[y_c, 1 : Nx + 1]) / np.concatenate([a_x[y_c, 0 : nw_half], np.ones(x_c)])
# bc_t_A1(1:nw_half) = C1_x(y_c,1:nw_half).*R_x2_c(y_c,1:nw_half);
        bc_t_A1[:nw_half] = (C1_x[y_c, :nw_half] * R_x2_c[y_c, :nw_half])
# bc_t_A1(Nx+2:Nx+x_c) = C3_x(y_c,1:nw_half).*R_x2_c(y_c,1:nw_half);
        bc_t_A1[Nx + 1: Nx + x_c + 1] = (C3_x[y_c, :nw_half] * R_x2_c[y_c, :nw_half])
# bc_b_A1(x_c:Nx) = C1_x(y_c,x_c+1:Nx+1).*R_x2_c(y_c,x_c+1:Nx+1);
        bc_b_A1[x_c : Nx] = (C1_x[y_c, x_c + 1 : Nx + 1] * R_x2_c[y_c, x_c + 1 : Nx + 1])
# bc_b_A1(Nx+x_c+1:2*Nx+1) = C3_x(y_c,x_c+1:Nx+1).*R_x2_c(y_c,x_c+1:Nx+1);
        bc_b_A1[Nx + x_c + 1 : 2 * Nx + 1] = (C3_x[y_c, x_c + 1 : Nx + 1] * R_x2_c[y_c, x_c + 1 : Nx + 1])
# d_t_xy_A1(1:nw_half) = -C2_x(y_c,1:nw_half).*R_x2_c(y_c,1:nw_half);
        d_t_xy_A1[:nw_half] = (-C2_x[y_c, :nw_half] * R_x2_c[y_c, :nw_half])
# d_b_xy_A1(x_c+1:Nx+1) = -C2_x(y_c,x_c+1:Nx+1).*R_x2_c(y_c,x_c+1:Nx+1);
        d_b_xy_A1[x_c + 1: Nx + 1] = (-C2_x[y_c, x_c + 1 : Nx + 1] * R_x2_c[y_c, x_c + 1 : Nx + 1])
# d_t_yx_A1(2:x_c) = -C4_x(y_c,1:nw_half).*R_x2_c(y_c,1:nw_half);
        d_t_yx_A1[1:x_c + 1] = (-C4_x[y_c, :nw_half] * R_x2_c[y_c, :nw_half])
# d_b_yx_A1(x_c:Nx) = -C4_x(y_c,x_c+1:Nx+1).*R_x2_c(y_c,x_c+1:Nx+1);
        d_b_yx_A1[x_c : Nx] = (-C4_x[y_c, x_c + 1 : Nx + 1] * R_x2_c[y_c, x_c + 1 : Nx + 1])
        A1_c = build_block_diag_A(
            d_c_A1[1 : Nx + 2],
            d_t_A1,
            d_b_A1,
            bc_t_A1,
            bc_b_A1,
            d_t_xy_A1,
            d_b_xy_A1,
            d_b_yx_A1,
            d_t_yx_A1,
            Nx,
        )
        # 右端项 (按列构成)
        for j in range(2, Nx + 1):
            bx_1[j - 1, 0] = cy * Wx[y_c - 1, j - 1] + (1 - 2 * cy) * Wx[y_c, j - 1] + cy * Wx[y_c + 1, j - 1]
            by_1[j - 1, 0] = cy * Wy[y_c - 1, j - 1] + (1 - 2 * cy) * Wy[y_c, j - 1] + cy * Wy[y_c + 1, j - 1]
            
        rhs = np.concatenate([
            (bx_1[:, 0] * R_in_d[y_c, :]).astype(complex),
            (by_1[:, 0] * R_in_d[y_c, :]).astype(complex),
        ])
        w1_c = np.linalg.solve(A1_c, rhs)
        Wx_m[y_c,  : ] = w1_c[: Nx + 1]
        Wy_m[y_c,  : ] = w1_c[Nx + 1 :]

        # ---------- A1 上下成对行 ----------
        timeCount(f'in z for 2')
        for i in range(1, min(nh_bottom, nh_top)):
            y_u = y_c - i
            y_d = y_c + i
            
            # 上行
            d_c_A1_u = 1 + (2 * cx) / a_x[y_u, :]
            d_t_A1_u = (-2 * cx) / (1 + a_x[y_u, : Nx ]) / np.concatenate([np.ones(x_c + 1), a_x[y_u, x_c + 1 : Nx]])
            d_b_A1_u = (-2 * cx) / (1 + a_x[y_u, 1 : Nx + 1]) / np.concatenate([a_x[y_u, 1 : nw_half], np.ones(x_c + 1)])
            
            bc_t_A1_u[:nw_half] = (C1_x[y_u, :nw_half] * R_x2_c[y_u, :nw_half])
            bc_t_A1_u[Nx + 1 : Nx + x_c + 1] = (C3_x[y_u, :nw_half] * R_x2_c[y_u, :nw_half])
            bc_b_A1_u[x_c : Nx] = (C1_x[y_u, x_c + 1 : Nx + 1] * R_x2_c[y_u, x_c + 1 : Nx + 1])
            bc_b_A1_u[Nx + x_c + 1: 2 * Nx + 1] = (C3_x[y_u, x_c + 1 : Nx + 1] * R_x2_c[y_u, x_c + 1 : Nx + 1])
            
            d_t_xy_A1_u[:nw_half] = (-C2_x[y_u, :nw_half] * R_x2_c[y_u, :nw_half])
            d_b_xy_A1_u[x_c + 1 : Nx + 1] = (-C2_x[y_u, x_c + 1 : Nx + 1] * R_x2_c[y_u, x_c + 1 : Nx + 1])
            d_t_yx_A1_u[1:x_c + 1] = (-C4_x[y_u, :nw_half] * R_x2_c[y_u, :nw_half])
            d_b_yx_A1_u[x_c : Nx] = (-C4_x[y_u, x_c + 1 : Nx + 1] * R_x2_c[y_u, x_c + 1 : Nx + 1])
            
            
            # 下行
            d_c_A1_d = 1 + (2 * cx) / a_x[y_d, :]
            d_t_A1_d = (-2 * cx) / (1 + a_x[y_d,  : Nx]) / np.concatenate([np.ones(x_c + 1), a_x[y_d, x_c + 1 : Nx]])
            d_b_A1_d = (-2 * cx) / (1 + a_x[y_d, 1 : Nx + 1]) / np.concatenate([a_x[y_d, 1 : nw_half], np.ones(x_c + 1)])
            
            bc_t_A1_d[:nw_half] = (C1_x[y_d, :nw_half] * R_x2_c[y_d, :nw_half])
            bc_t_A1_d[Nx + 1 : Nx + x_c + 1] = (C3_x[y_d, :nw_half] * R_x2_c[y_d, :nw_half])
            bc_b_A1_d[x_c : Nx] = (C1_x[y_d, x_c + 1 : Nx + 1] * R_x2_c[y_d, x_c + 1 : Nx + 1])
            bc_b_A1_d[Nx + x_c + 1 : 2 * Nx + 1] = (C3_x[y_d, x_c + 1 : Nx + 1] * R_x2_c[y_d, x_c + 1 : Nx + 1])
            
            d_t_xy_A1_d[:nw_half] = (-C2_x[y_d, :nw_half] * R_x2_c[y_d, :nw_half])
            d_b_xy_A1_d[x_c + 1 : Nx + 1] = (-C2_x[y_d, x_c + 1 : Nx + 1] * R_x2_c[y_d, x_c + 1 : Nx + 1])
            d_t_yx_A1_d[1:x_c + 1] = (-C4_x[y_d, :nw_half] * R_x2_c[y_d, :nw_half])
            d_b_yx_A1_d[x_c: Nx] = (-C4_x[y_d, x_c + 1 : Nx + 1] * R_x2_c[y_d, x_c + 1 : Nx + 1])
            
            A1_u = build_block_diag_A(
                d_c_A1_u[1 : Nx + 2],
                d_t_A1_u,
                d_b_A1_u,
                bc_t_A1_u,
                bc_b_A1_u,
                d_t_xy_A1_u,
                d_b_xy_A1_u,
                d_b_yx_A1_u,
                d_t_yx_A1_u,
                Nx,
            )
            
            A1_d = build_block_diag_A(
                d_c_A1_d[1 : Nx + 2],
                d_t_A1_d,
                d_b_A1_d,
                bc_t_A1_d,
                bc_b_A1_d,
                d_t_xy_A1_d,
                d_b_xy_A1_d,
                d_b_yx_A1_d,
                d_t_yx_A1_d,
                Nx,
            )
            # 右端构造 (省略若干中间赋值，严格对应MATLAB原式)
            j = x_c
            bx_1_u[j, 0] = ( (2 * cy) / (1 + a_y[y_u, j]) / a_y[y_u, j] * Wx[y_u - 1, j] + (1 - (2 * cy) / a_y[y_u, j]) * Wx[y_u, j] * R_in_d[y_u, j] + (2 * cy) / (1 + a_y[y_u, j]) * Wx[y_u + 1, j] + R_x2_n[y_u, j] * C1_x[y_u, j] * Wx_m[y_u + 1, j] + R_x2_n[y_u, j] * C2_x[y_u, j] * Wy_m[y_u + 1, j])
            by_1_u[j - 1, 0] = ( (2 * cy) / (1 + a_y[y_u, j]) / a_y[y_u, j] * Wy[y_u - 1, j] + (1 - (2 * cy) / a_y[y_u, j]) * Wy[y_u, j] * R_in_d[y_u, j] + (2 * cy) / (1 + a_y[y_u, j]) * Wy[y_u + 1, j] + R_x2_n[y_u, j] * C3_x[y_u, j] * Wy_m[y_u + 1, j] + R_x2_n[y_u, j] * C4_x[y_u, j] * Wx_m[y_u + 1, j])
            bx_1_d[j - 1, 0] = ( (2 * cy) / (1 + a_y[y_d, j]) * Wx[y_d - 1, j] + (1 - (2 * cy) / a_y[y_d, j]) * Wx[y_d, j] * R_in_d[y_d, j] + (2 * cy) / (1 + a_y[y_d, j]) / a_y[y_d, j] * Wx[y_d + 1, j] + R_x2_n[y_d, j] * C1_x[y_d, j] * Wx_m[y_d - 1, j] + R_x2_n[y_d, j] * C2_x[y_d, j] * Wy_m[y_d - 1, j])
            by_1_d[j - 1, 0] = ( (2 * cy) / (1 + a_y[y_d, j]) * Wy[y_d - 1, j] + (1 - (2 * cy) / a_y[y_d, j]) * Wy[y_d, j] * R_in_d[y_d, j] + (2 * cy) / (1 + a_y[y_d, j]) / a_y[y_d, j] * Wy[y_d + 1, j] + R_x2_n[y_d, j] * C3_x[y_d, j] * Wy_m[y_d - 1, j] + R_x2_n[y_d, j] * C4_x[y_d, j] * Wx_m[y_d - 1, j])
            timeCount(f'in z for 2-1, {i}')
            for j in range(1, nw_half + 1):
                xx_l = x_c - j
                xx_r = x_c + j
                # 上行右侧
                bx_1_u[xx_r, 0] = ( (2 * cy) / (1 + a_y[y_u, xx_r]) / a_y[y_u, xx_r] * Wx[y_u - 1, xx_r] + (1 - (2 * cy) / a_y[y_u, xx_r]) * Wx[y_u, xx_r] * R_in_d[y_u, xx_r] + (2 * cy) / (1 + a_y[y_u, xx_r])                            * Wx[y_u + 1, xx_r] + R_x2[y_u, xx_r] * C1_x[y_u, xx_r] * Wx_m[y_u + 1, xx_r - 1] + R_x2_n[y_u, xx_r] * C1_x[y_u, xx_r] * Wx_m[y_u + 1, xx_r] + R_x2[y_u, xx_r] * C2_x[y_u, xx_r] * Wy_m[y_u + 1, xx_r - 1] + R_x2_n[y_u, xx_r] * C2_x[y_u, xx_r] * Wy_m[y_u + 1, xx_r])
                by_1_u[xx_r, 0] = ( (2 * cy) / (1 + a_y[y_u, xx_r]) / a_y[y_u, xx_r] * Wy[y_u - 1, xx_r] + (1 - (2 * cy) / a_y[y_u, xx_r]) * Wy[y_u, xx_r] * R_in_d[y_u, xx_r] + (2 * cy) / (1 + a_y[y_u, xx_r])                            * Wy[y_u + 1, xx_r] + R_x2[y_u, xx_r] * C3_x[y_u, xx_r] * Wy_m[y_u + 1, xx_r - 1] + R_x2_n[y_u, xx_r] * C3_x[y_u, xx_r] * Wy_m[y_u + 1, xx_r] + R_x2[y_u, xx_r] * C4_x[y_u, xx_r] * Wx_m[y_u + 1, xx_r - 1] + R_x2_n[y_u, xx_r] * C4_x[y_u, xx_r] * Wx_m[y_u + 1, xx_r])
                # 上行左侧
                bx_1_u[xx_l, 0] = ( (2 * cy) / (1 + a_y[y_u, xx_l]) / a_y[y_u, xx_l] * Wx[y_u - 1, xx_l]   + (1 - (2 * cy) / a_y[y_u, xx_l]) * Wx[y_u, xx_l] * R_in_d[y_u, xx_l]    + (2 * cy) / (1 + a_y[y_u, xx_l])                            * Wx[y_u + 1, xx_l] + R_x2[y_u, xx_l] * C1_x[y_u, xx_l] * Wx_m[y_u + 1, xx_l + 1]    + R_x2_n[y_u, xx_l] * C1_x[y_u, xx_l]  * Wx_m[y_u + 1, xx_l] + R_x2[y_u, xx_l] * C2_x[y_u, xx_l] * Wy_m[y_u + 1, xx_l + 1]     + R_x2_n[y_u, xx_l] * C2_x[y_u, xx_l] * Wy_m[y_u + 1, xx_l])
                by_1_u[xx_l, 0] = ( (2 * cy) / (1 + a_y[y_u, xx_l]) / a_y[y_u, xx_l] * Wy[y_u - 1, xx_l]   + (1 - (2 * cy) / a_y[y_u, xx_l]) * Wy[y_u, xx_l] * R_in_d[y_u, xx_l]    + (2 * cy) / (1 + a_y[y_u, xx_l])                            * Wy[y_u + 1, xx_l] + R_x2[y_u, xx_l] * C3_x[y_u, xx_l] * Wy_m[y_u + 1, xx_l + 1]    + R_x2_n[y_u, xx_l] * C3_x[y_u, xx_l]  * Wy_m[y_u + 1, xx_l] + R_x2[y_u, xx_l] * C4_x[y_u, xx_l] * Wx_m[y_u + 1, xx_l + 1]     + R_x2_n[y_u, xx_l] * C4_x[y_u, xx_l] * Wx_m[y_u + 1, xx_l])
                # 下行右侧
                bx_1_d[xx_r, 0] = ( (2 * cy) / (1 + a_y[y_d, xx_r])                          * Wx[y_d - 1, xx_r] + (1 - (2 * cy) / a_y[y_d, xx_r]) * Wx[y_d, xx_r] * R_in_d[y_d, xx_r] + (2 * cy) / (1 + a_y[y_d, xx_r]) / a_y[y_d, xx_r] * Wx[y_d + 1, xx_r] + R_x2[y_d, xx_r] * C1_x[y_d, xx_r] * Wx_m[y_d - 1, xx_r - 1] + R_x2_n[y_d, xx_r] * C1_x[y_d, xx_r] * Wx_m[y_d - 1, xx_r] + R_x2[y_d, xx_r] * C2_x[y_d, xx_r] * Wy_m[y_d - 1, xx_r - 1] + R_x2_n[y_d, xx_r] * C2_x[y_d, xx_r] * Wy_m[y_d - 1, xx_r])
                by_1_d[xx_r, 0] = ( (2 * cy) / (1 + a_y[y_d, xx_r])                          * Wy[y_d - 1, xx_r] + (1 - (2 * cy) / a_y[y_d, xx_r]) * Wy[y_d, xx_r] * R_in_d[y_d, xx_r] + (2 * cy) / (1 + a_y[y_d, xx_r]) / a_y[y_d, xx_r] * Wy[y_d + 1, xx_r] + R_x2[y_d, xx_r] * C3_x[y_d, xx_r] * Wy_m[y_d - 1, xx_r - 1] + R_x2_n[y_d, xx_r] * C3_x[y_d, xx_r] * Wy_m[y_d - 1, xx_r] + R_x2[y_d, xx_r] * C4_x[y_d, xx_r] * Wx_m[y_d - 1, xx_r - 1] + R_x2_n[y_d, xx_r] * C4_x[y_d, xx_r] * Wx_m[y_d - 1, xx_r])
                # 下行左侧
                bx_1_d[xx_l, 0] = ( (2 * cy) / (1 + a_y[y_d, xx_l])                           * Wx[y_d - 1, xx_l] + (1 - (2 * cy) / a_y[y_d, xx_l]) * Wx[y_d, xx_l] * R_in_d[y_d, xx_l]    + (2 * cy) / (1 + a_y[y_d, xx_l]) / a_y[y_d, xx_l] * Wx[y_d + 1, xx_l] + R_x2[y_d, xx_l] * C1_x[y_d, xx_l] * Wx_m[y_d - 1, xx_l + 1]    + R_x2_n[y_d, xx_l] * C1_x[y_d, xx_l]   * Wx_m[y_d - 1, xx_l] + R_x2[y_d, xx_l] * C2_x[y_d, xx_l] * Wy_m[y_d - 1, xx_l + 1]     + R_x2_n[y_d, xx_l] * C2_x[y_d, xx_l] * Wy_m[y_d - 1, xx_l])
                by_1_d[xx_l, 0] = ( (2 * cy) / (1 + a_y[y_d, xx_l])                           * Wy[y_d - 1, xx_l] + (1 - (2 * cy) / a_y[y_d, xx_l]) * Wy[y_d, xx_l] * R_in_d[y_d, xx_l]    + (2 * cy) / (1 + a_y[y_d, xx_l]) / a_y[y_d, xx_l] * Wy[y_d + 1, xx_l] + R_x2[y_d, xx_l] * C3_x[y_d, xx_l] * Wy_m[y_d - 1, xx_l + 1]    + R_x2_n[y_d, xx_l] * C3_x[y_d, xx_l]   * Wy_m[y_d - 1, xx_l] + R_x2[y_d, xx_l] * C4_x[y_d, xx_l] * Wx_m[y_d - 1, xx_l + 1]     + R_x2_n[y_d, xx_l] * C4_x[y_d, xx_l] * Wx_m[y_d - 1, xx_l])
            
            timeCount(f'in z for 2-2, {i}')                
            w1_u = np.linalg.solve( A1_u, np.concatenate([bx_1_u[:, 0], by_1_u[:, 0]]).astype(complex),)
            w1_d = np.linalg.solve( A1_d, np.concatenate([bx_1_d[:, 0], by_1_d[:, 0]]).astype(complex),)
            Wx_m[y_u, :] = w1_u[: Nx + 1]
            Wx_m[y_d, :] = w1_d[: Nx + 1]
            Wy_m[y_u, :] = w1_u[Nx + 1 :]
            Wy_m[y_d, :] = w1_d[Nx + 1 :]
            timeCount(f'in z for 2-3, {i}')

        # 上剩余部分
        timeCount(f'in z for 3')
        for i in range(nh_bottom, nh_top):
            y_u = y_c - i
            d_c_A1_u = 1 + (2 * cx) / a_x[y_u, :]
            d_t_A1_u = (-2 * cx) / (1 + a_x[y_u, 1 : Nx + 1]) / np.concatenate([np.ones(1), a_x[y_u, x_c + 1 : Nx]])
            d_b_A1_u = (-2 * cx) / (1 + a_x[y_u, 2 : Nx + 2]) / np.concatenate([a_x[y_u, 2 : nw_half + 1], np.ones(1)])
            bc_t_A1_u[:nw_half] = (C1_x[y_u, :nw_half] * R_x2_c[y_u, :nw_half])
            bc_t_A1_u[Nx + 1 + 1 : Nx + x_c + 1] = (C3_x[y_u, :nw_half] * R_x2_c[y_u, :nw_half])
            bc_b_A1_u[x_c - 1 : Nx] = (C1_x[y_u, x_c + 1 : Nx + 2] * R_x2_c[y_u, x_c + 1 : Nx + 2])
            bc_b_A1_u[Nx + x_c : 2 * Nx + 1] = (C3_x[y_u, x_c + 1 : Nx + 2] * R_x2_c[y_u, x_c + 1 : Nx + 2])
            d_t_xy_A1_u[:nw_half] = (-C2_x[y_u, :nw_half] * R_x2_c[y_u, :nw_half])
            d_b_xy_A1_u[x_c : Nx + 2] = (-C2_x[y_u, x_c + 1 : Nx + 2] * R_x2_c[y_u, x_c + 1 : Nx + 2])
            d_t_yx_A1_u[1:x_c + 1] = (-C4_x[y_u, :nw_half] * R_x2_c[y_u, :nw_half])
            d_b_yx_A1_u[x_c - 1 : Nx] = (-C4_x[y_u, x_c + 1 : Nx + 2] * R_x2_c[y_u, x_c + 1 : Nx + 2])
            A1_u = build_block_diag_A(
                d_c_A1_u[1 : Nx + 2],
                d_t_A1_u,
                d_b_A1_u,
                bc_t_A1_u,
                bc_b_A1_u,
                d_t_xy_A1_u,
                d_b_xy_A1_u,
                d_b_yx_A1_u,
                d_t_yx_A1_u,
                Nx,
            )
            j = x_c
            bx_1_u[j - 1, 0] = ( (2 * cy) / (1 + a_y[y_u, j]) / a_y[y_u, j] * Wx[y_u - 1, j] + (1 - (2 * cy) / a_y[y_u, j]) * Wx[y_u, j] * R_in_d[y_u, j] + (2 * cy) / (1 + a_y[y_u, j]) * Wx[y_u + 1, j] + R_x2_n[y_u, j] * C1_x[y_u, j] * Wx_m[y_u + 1, j] + R_x2_n[y_u, j] * C2_x[y_u, j] * Wy_m[y_u + 1, j])
            by_1_u[j - 1, 0] = ( (2 * cy) / (1 + a_y[y_u, j]) / a_y[y_u, j] * Wy[y_u - 1, j] + (1 - (2 * cy) / a_y[y_u, j]) * Wy[y_u, j] * R_in_d[y_u, j] + (2 * cy) / (1 + a_y[y_u, j]) * Wy[y_u + 1, j] + R_x2_n[y_u, j] * C3_x[y_u, j] * Wy_m[y_u + 1, j] + R_x2_n[y_u, j] * C4_x[y_u, j] * Wx_m[y_u + 1, j])
            for j in range(1, nw_half + 1):
                xx_l = x_c - j
                xx_r = x_c + j
                # 同上 (仅上行)
                bx_1_u[xx_r - 1, 0] = ( (2 * cy) / (1 + a_y[y_u, xx_r]) / a_y[y_u, xx_r] * Wx[y_u - 1, xx_r] + (1 - (2 * cy) / a_y[y_u, xx_r]) * Wx[y_u, xx_r] * R_in_d[y_u, xx_r] + (2 * cy) / (1 + a_y[y_u, xx_r]) * Wx[y_u + 1, xx_r] + R_x2[y_u, xx_r] * C1_x[y_u, xx_r] * Wx_m[y_u + 1, xx_r - 1] + R_x2_n[y_u, xx_r] * C1_x[y_u, xx_r] * Wx_m[y_u + 1, xx_r] + R_x2[y_u, xx_r] * C2_x[y_u, xx_r] * Wy_m[y_u + 1, xx_r - 1] + R_x2_n[y_u, xx_r] * C2_x[y_u, xx_r] * Wy_m[y_u + 1, xx_r])
                by_1_u[xx_r - 1, 0] = ( (2 * cy) / (1 + a_y[y_u, xx_r]) / a_y[y_u, xx_r] * Wy[y_u - 1, xx_r] + (1 - (2 * cy) / a_y[y_u, xx_r]) * Wy[y_u, xx_r] * R_in_d[y_u, xx_r] + (2 * cy) / (1 + a_y[y_u, xx_r]) * Wy[y_u + 1, xx_r] + R_x2[y_u, xx_r] * C3_x[y_u, xx_r] * Wy_m[y_u + 1, xx_r - 1] + R_x2_n[y_u, xx_r] * C3_x[y_u, xx_r] * Wy_m[y_u + 1, xx_r] + R_x2[y_u, xx_r] * C4_x[y_u, xx_r] * Wx_m[y_u + 1, xx_r - 1] + R_x2_n[y_u, xx_r] * C4_x[y_u, xx_r] * Wx_m[y_u + 1, xx_r])
                bx_1_u[xx_l - 1, 0] = ( (2 * cy) / (1 + a_y[y_u, xx_l]) / a_y[y_u, xx_l] * Wx[y_u - 1, xx_l] + (1 - (2 * cy) / a_y[y_u, xx_l]) * Wx[y_u, xx_l] * R_in_d[y_u, xx_l] + (2 * cy) / (1 + a_y[y_u, xx_l]) * Wx[y_u + 1, xx_l] + R_x2[y_u, xx_l] * C1_x[y_u, xx_l] * Wx_m[y_u + 1, xx_l + 1] + R_x2_n[y_u, xx_l] * C1_x[y_u, xx_l] * Wx_m[y_u + 1, xx_l] + R_x2[y_u, xx_l] * C2_x[y_u, xx_l] * Wy_m[y_u + 1, xx_l + 1] + R_x2_n[y_u, xx_l] * C2_x[y_u, xx_l] * Wy_m[y_u + 1, xx_l])
                by_1_u[xx_l - 1, 0] = ( (2 * cy) / (1 + a_y[y_u, xx_l]) / a_y[y_u, xx_l] * Wy[y_u - 1, xx_l] + (1 - (2 * cy) / a_y[y_u, xx_l]) * Wy[y_u, xx_l] * R_in_d[y_u, xx_l] + (2 * cy) / (1 + a_y[y_u, xx_l]) * Wy[y_u + 1, xx_l] + R_x2[y_u, xx_l] * C3_x[y_u, xx_l] * Wy_m[y_u + 1, xx_l + 1] + R_x2_n[y_u, xx_l] * C3_x[y_u, xx_l] * Wy_m[y_u + 1, xx_l] + R_x2[y_u, xx_l] * C4_x[y_u, xx_l] * Wx_m[y_u + 1, xx_l + 1] + R_x2_n[y_u, xx_l] * C4_x[y_u, xx_l] * Wx_m[y_u + 1, xx_l])
            w1_u = np.linalg.solve( A1_u, np.concatenate([bx_1_u[:, 0], by_1_u[:, 0]]).astype(complex))
            Wx_m[y_u, 1 : Nx + 2] = w1_u[: Nx + 1]
            Wy_m[y_u, 1 : Nx + 2] = w1_u[Nx + 1 :]

        # 下剩余部分
        timeCount(f'in z for 4')
        for i in range(nh_top, nh_bottom):
            y_d = y_c + i
            d_c_A1_d = 1 + (2 * cx) / a_x[y_d, :]
            d_t_A1_d = (-2 * cx) / (1 + a_x[y_d, 1 : Nx + 1]) / np.concatenate([np.ones(1), a_x[y_d, x_c + 1 : Nx]])
            d_b_A1_d = (-2 * cx) / (1 + a_x[y_d, 2 : Nx + 2]) / np.concatenate([a_x[y_d, 2 : nw_half + 1], np.ones(1)])
            bc_t_A1_d[:nw_half] = (C1_x[y_d, :nw_half] * R_x2_c[y_d, :nw_half])
            bc_t_A1_d[Nx + 1 + 1 : Nx + x_c + 1] = (C3_x[y_d, :nw_half] * R_x2_c[y_d, :nw_half])
            bc_b_A1_d[x_c - 1 : Nx] = (C1_x[y_d, x_c + 1 : Nx + 2] * R_x2_c[y_d, x_c + 1 : Nx + 2])
            bc_b_A1_d[Nx + x_c : 2 * Nx + 1] = (C3_x[y_d, x_c + 1 : Nx + 2] * R_x2_c[y_d, x_c + 1 : Nx + 2])
            d_t_xy_A1_d[:nw_half] = (-C2_x[y_d, :nw_half] * R_x2_c[y_d, :nw_half])
            d_b_xy_A1_d[x_c : Nx + 2] = (-C2_x[y_d, x_c + 1 : Nx + 2] * R_x2_c[y_d, x_c + 1 : Nx + 2])
            d_t_yx_A1_d[1:x_c + 1] = (-C4_x[y_d, :nw_half] * R_x2_c[y_d, :nw_half])
            d_b_yx_A1_d[x_c - 1 : Nx] = (-C4_x[y_d, x_c + 1 : Nx + 2] * R_x2_c[y_d, x_c + 1 : Nx + 2])
            A1_d = build_block_diag_A(
                d_c_A1_d[1 : Nx + 2],
                d_t_A1_d,
                d_b_A1_d,
                bc_t_A1_d,
                bc_b_A1_d,
                d_t_xy_A1_d,
                d_b_xy_A1_d,
                d_b_yx_A1_d,
                d_t_yx_A1_d,
                Nx,
            )
            j = x_c
            bx_1_d[j - 1, 0] = ( (2 * cy) / (1 + a_y[y_d, j]) * Wx[y_d - 1, j] + (1 - (2 * cy) / a_y[y_d, j]) * Wx[y_d, j] * R_in_d[y_d, j] + (2 * cy) / (1 + a_y[y_d, j]) / a_y[y_d, j] * Wx[y_d + 1, j] + R_x2_n[y_d, j] * C1_x[y_d, j] * Wx_m[y_d - 1, j] + R_x2_n[y_d, j] * C2_x[y_d, j] * Wy_m[y_d - 1, j])
            by_1_d[j - 1, 0] = ( (2 * cy) / (1 + a_y[y_d, j]) * Wy[y_d - 1, j] + (1 - (2 * cy) / a_y[y_d, j]) * Wy[y_d, j] * R_in_d[y_d, j] + (2 * cy) / (1 + a_y[y_d, j]) / a_y[y_d, j] * Wy[y_d + 1, j] + R_x2_n[y_d, j] * C3_x[y_d, j] * Wy_m[y_d - 1, j] + R_x2_n[y_d, j] * C4_x[y_d, j] * Wx_m[y_d - 1, j])
            for j in range(1, nw_half + 1):
                xx_l = x_c - j
                xx_r = x_c + j
                # 同上 (仅下行)
                bx_1_d[xx_r - 1, 0] = ( (2 * cy) / (1 + a_y[y_d, xx_r]) * Wx[y_d - 1, xx_r] + (1 - (2 * cy) / a_y[y_d, xx_r]) * Wx[y_d, xx_r] * R_in_d[y_d, xx_r] + (2 * cy) / (1 + a_y[y_d, xx_r]) / a_y[y_d, xx_r] * Wx[y_d + 1, xx_r] + R_x2[y_d, xx_r] * C1_x[y_d, xx_r] * Wx_m[y_d - 1, xx_r - 1] + R_x2_n[y_d, xx_r] * C1_x[y_d, xx_r] * Wx_m[y_d - 1, xx_r] + R_x2[y_d, xx_r] * C2_x[y_d, xx_r] * Wy_m[y_d - 1, xx_r - 1] + R_x2_n[y_d, xx_r] * C2_x[y_d, xx_r] * Wy_m[y_d - 1, xx_r])
                by_1_d[xx_r - 1, 0] = ( (2 * cy) / (1 + a_y[y_d, xx_r]) * Wy[y_d - 1, xx_r] + (1 - (2 * cy) / a_y[y_d, xx_r]) * Wy[y_d, xx_r] * R_in_d[y_d, xx_r] + (2 * cy) / (1 + a_y[y_d, xx_r]) / a_y[y_d, xx_r] * Wy[y_d + 1, xx_r] + R_x2[y_d, xx_r] * C3_x[y_d, xx_r] * Wy_m[y_d - 1, xx_r - 1] + R_x2_n[y_d, xx_r] * C3_x[y_d, xx_r] * Wy_m[y_d - 1, xx_r] + R_x2[y_d, xx_r] * C4_x[y_d, xx_r] * Wx_m[y_d - 1, xx_r - 1] + R_x2_n[y_d, xx_r] * C4_x[y_d, xx_r] * Wx_m[y_d - 1, xx_r])
                bx_1_d[xx_l - 1, 0] = ( (2 * cy) / (1 + a_y[y_d, xx_l]) * Wx[y_d - 1, xx_l] + (1 - (2 * cy) / a_y[y_d, xx_l]) * Wx[y_d, xx_l] * R_in_d[y_d, xx_l] + (2 * cy) / (1 + a_y[y_d, xx_l]) / a_y[y_d, xx_l] * Wx[y_d + 1, xx_l] + R_x2[y_d, xx_l] * C1_x[y_d, xx_l] * Wx_m[y_d - 1, xx_l + 1] + R_x2_n[y_d, xx_l] * C1_x[y_d, xx_l] * Wx_m[y_d - 1, xx_l] + R_x2[y_d, xx_l] * C2_x[y_d, xx_l] * Wy_m[y_d - 1, xx_l + 1] + R_x2_n[y_d, xx_l] * C2_x[y_d, xx_l] * Wy_m[y_d - 1, xx_l])
                by_1_d[xx_l - 1, 0] = ( (2 * cy) / (1 + a_y[y_d, xx_l]) * Wy[y_d - 1, xx_l] + (1 - (2 * cy) / a_y[y_d, xx_l]) * Wy[y_d, xx_l] * R_in_d[y_d, xx_l] + (2 * cy) / (1 + a_y[y_d, xx_l]) / a_y[y_d, xx_l] * Wy[y_d + 1, xx_l] + R_x2[y_d, xx_l] * C3_x[y_d, xx_l] * Wy_m[y_d - 1, xx_l + 1] + R_x2_n[y_d, xx_l] * C3_x[y_d, xx_l] * Wy_m[y_d - 1, xx_l] + R_x2[y_d, xx_l] * C4_x[y_d, xx_l] * Wx_m[y_d - 1, xx_l + 1] + R_x2_n[y_d, xx_l] * C4_x[y_d, xx_l] * Wx_m[y_d - 1, xx_l])
            w1_d = np.linalg.solve(
                A1_d, np.concatenate([bx_1_d[:, 0], by_1_d[:, 0]]).astype(complex)
            )
            Wx_m[y_d, 1 : Nx + 2] = w1_d[: Nx + 1]
            Wy_m[y_d, 1 : Nx + 2] = w1_d[Nx + 1 :]

        # ---------- A2 中央列 ----------
        
    # d_c_A2 = 1 + (2*cy)./a_y(:,x_c);
    # d_t_A2 = (-2*cy)./(1+a_y(1:Ny,x_c))./([ones(y_c,1); a_y(y_c+1:Ny,x_c)]);
    # d_b_A2 = (-2*cy)./(1+a_y(2:Ny+1,x_c))./([a_y(2:nh_top,x_c); ones(nh_bottom+1,1)]);
    # bc_t_A2(1:nh_top) = C1_y(1:nh_top,x_c).*R_y1_c(1:nh_top,x_c);
    # bc_t_A2(Ny+2:Ny+y_c) = C3_y(1:nh_top,x_c).*R_y1_c(1:nh_top,x_c);
    # bc_b_A2(y_c:Ny) = C1_y(y_c+1:Ny+1,x_c).*R_y1_c(y_c+1:Ny+1,x_c);
    # bc_b_A2(Ny+y_c+1:2*Ny+1) = C3_y(y_c+1:Ny+1,x_c).*R_y1_c(y_c+1:Ny+1,x_c);
    # d_t_xy_A2(1:nh_top) = -C2_y(1:nh_top,x_c).*R_y1_c(1:nh_top,x_c);
    # d_b_xy_A2(y_c+1:Ny+1) = -C2_y(y_c+1:Ny+1,x_c).*R_y1_c(y_c+1:Ny+1,x_c);
    # d_t_yx_A2(2:y_c) = -C4_y(1:nh_top,x_c).*R_y1_c(1:nh_top,x_c);
    # d_b_yx_A2(y_c:Ny) = -C4_y(y_c+1:Ny+1,x_c).*R_y1_c(y_c+1:Ny+1,x_c);
        d_c_A2 = 1 + (2 * cy) / a_y[:, x_c]
        d_t_A2 = (-2 * cy) / (1 + a_y[:Ny, x_c]) / np.concatenate([np.ones(y_c + 1), a_y[y_c + 1 : Ny, x_c]])
        d_b_A2 = (-2 * cy) / (1 + a_y[1 : Ny + 1, x_c]) / np.concatenate([a_y[1:nh_top, x_c], np.ones(nh_bottom + 1)])
        bc_t_A2[:nh_top, 0] = (C1_y[:nh_top, x_c] * R_y1_c[:nh_top, x_c])
        bc_t_A2[Ny + 1 : Ny + y_c + 1, 0] = (C3_y[:nh_top, x_c] * R_y1_c[:nh_top, x_c])
        bc_b_A2[y_c : Ny, 0] = (C1_y[y_c + 1 : Ny + 1, x_c] * R_y1_c[y_c + 1 : Ny + 1, x_c])
        bc_b_A2[Ny + y_c + 1 : 2 * Ny + 1, 0] = (C3_y[y_c + 1 : Ny + 1, x_c] * R_y1_c[y_c + 1 : Ny + 1, x_c])
        d_t_xy_A2[:nh_top, 0] = (-C2_y[:nh_top, x_c] * R_y1_c[:nh_top, x_c])
        d_b_xy_A2[y_c + 1 : Ny + 1, 0] = (-C2_y[y_c + 1 : Ny + 1, x_c] * R_y1_c[y_c + 1 : Ny + 1, x_c])
        d_t_yx_A2[1:y_c + 1, 0] = (-C4_y[:nh_top, x_c] * R_y1_c[:nh_top, x_c])
        d_b_yx_A2[y_c : Ny, 0] = (-C4_y[y_c + 1 : Ny + 1, x_c] * R_y1_c[y_c + 1 : Ny + 1, x_c])
        A2_c = build_block_diag_A(
            d_c_A2.flatten(),
            d_t_A2,
            d_b_A2,
            bc_t_A2.flatten(),
            bc_b_A2.flatten(),
            d_t_xy_A2.flatten(),
            d_b_xy_A2.flatten(),
            d_b_yx_A2.flatten(),
            d_t_yx_A2.flatten(),
            Ny,
        )
        for ii in range(1, Ny):
            bx_2[ii, 0] = cx * Wx_m[ii, x_c - 1] + (1 - 2 * cx) * Wx_m[ii, x_c] + cx * Wx_m[ii, x_c + 1]
            by_2[ii, 0] = cx * Wy_m[ii, x_c - 1] + (1 - 2 * cx) * Wy_m[ii, x_c] + cx * Wy_m[ii, x_c + 1]
            
        rhs2 = np.concatenate([
            (bx_2[:, 0] * R_in_d[ :, x_c]).astype(complex),
            (by_2[:, 0] * R_in_d[ :, x_c]).astype(complex),
        ])
        w2_c = np.linalg.solve(A2_c, rhs2)
        Wx[:, x_c] = w2_c[: Ny + 1]
        Wy[:, x_c] = w2_c[Ny + 1 :]
        #11111111

        timeCount(f'in z for 5')
        # ---------- A2 左右列 ----------
        for jj in range(1, nw_half):
            x_l = x_c - jj
            x_r = x_c + jj
            # 左列
            
        # d_c_A2_l = 1 + (2*cy)./a_y(:,x_l);
        # d_t_A2_l = (-2*cy)./(1+a_y(1:Ny,x_l))./([ones(y_c,1); a_y(y_c+1:Ny,x_l)]);
        # d_b_A2_l = (-2*cy)./(1+a_y(2:Ny+1,x_l))./([a_y(2:nh_top,x_l); ones(nh_bottom+1,1)]);
            d_c_A2_l = 1 + (2 * cy) / a_y[:, x_l]
            d_t_A2_l = (-2 * cy) / (1 + a_y[:Ny, x_l]) / np.concatenate([np.ones(y_c + 1), a_y[y_c + 1 : Ny, x_l]])
            d_b_A2_l = (-2 * cy) / (1 + a_y[1 : Ny + 1, x_l]) / np.concatenate([a_y[1:nh_top, x_l], np.ones(nh_bottom + 1)])
            
        # bc_t_A2_l(1:nh_top) = C1_y(1:nh_top,x_l).*R_y1_c(1:nh_top,x_l);
        # bc_t_A2_l(Ny+2:Ny+y_c) = C3_y(1:nh_top,x_l).*R_y1_c(1:nh_top,x_l);
        # bc_b_A2_l(y_c:Ny) = C1_y(y_c+1:Ny+1,x_l).*R_y1_c(y_c+1:Ny+1,x_l);
        # bc_b_A2_l(Ny+y_c+1:2*Ny+1) = C3_y(y_c+1:Ny+1,x_l).*R_y1_c(y_c+1:Ny+1,x_l);
            bc_t_A2_l[:nh_top, 0] = (C1_y[:nh_top, x_l] * R_y1_c[:nh_top, x_l])
            bc_t_A2_l[Ny + 1 : Ny + y_c + 1, 0] = (C3_y[:nh_top, x_l] * R_y1_c[:nh_top, x_l])
            bc_b_A2_l[y_c : Ny, 0] = (C1_y[y_c + 1 : Ny + 1, x_l] * R_y1_c[y_c + 1 : Ny + 1, x_l])
            bc_b_A2_l[Ny + y_c + 1 : 2 * Ny + 1, 0] = (C3_y[y_c + 1 : Ny + 1, x_l] * R_y1_c[y_c + 1 : Ny + 1, x_l])

        # d_t_xy_A2_l(1:nh_top) = -C2_y(1:nh_top,x_l).*R_y1_c(1:nh_top,x_l);
        # d_b_xy_A2_l(y_c+1:Ny+1) = -C2_y(y_c+1:Ny+1,x_l).*R_y1_c(y_c+1:Ny+1,x_l);
        # d_t_yx_A2_l(2:y_c) = -C4_y(1:nh_top,x_l).*R_y1_c(1:nh_top,x_l);
        # d_b_yx_A2_l(y_c:Ny) = -C4_y(y_c+1:Ny+1,x_l).*R_y1_c(y_c+1:Ny+1,x_l);
            d_t_xy_A2_l[:nh_top, 0] = (-C2_y[:nh_top, x_l] * R_y1_c[:nh_top, x_l])
            d_b_xy_A2_l[y_c + 1 : Ny + 1, 0] = (-C2_y[y_c + 1 : Ny + 1, x_l] * R_y1_c[y_c + 1 : Ny + 1, x_l])
            d_t_yx_A2_l[1:y_c + 1, 0] = (-C4_y[:nh_top, x_l] * R_y1_c[:nh_top, x_l])
            d_b_yx_A2_l[y_c : Ny, 0] = (-C4_y[y_c + 1 : Ny + 1, x_l] * R_y1_c[y_c + 1 : Ny + 1, x_l])

            # 右列
        # d_c_A2_r = 1 + (2*cy)./a_y(:,x_r);
        # d_t_A2_r = (-2*cy)./(1+a_y(1:Ny,x_r))./([ones(y_c,1); a_y(y_c+1:Ny,x_r)]);
        # d_b_A2_r = (-2*cy)./(1+a_y(2:Ny+1,x_r))./([a_y(2:nh_top,x_r); ones(nh_bottom+1,1)]);
            d_c_A2_r = 1 + (2 * cy) / a_y[:, x_r]
            d_t_A2_r = (-2 * cy) / (1 + a_y[:Ny, x_r]) / np.concatenate([np.ones(y_c  + 1), a_y[y_c + 1 : Ny, x_r]])
            d_b_A2_r = (-2 * cy) / (1 + a_y[1 : Ny + 1, x_r]) / np.concatenate([a_y[1:nh_top, x_r], np.ones(nh_bottom + 1)])
            
        # bc_t_A2_r(1:nh_top) = C1_y(1:nh_top,x_r).*R_y1_c(1:nh_top,x_r);
        # bc_t_A2_r(Ny+2:Ny+y_c) = C3_y(1:nh_top,x_r).*R_y1_c(1:nh_top,x_r);
        # bc_b_A2_r(y_c:Ny) = C1_y(y_c+1:Ny+1,x_r).*R_y1_c(y_c+1:Ny+1,x_r);
        # bc_b_A2_r(Ny+y_c+1:2*Ny+1) = C3_y(y_c+1:Ny+1,x_r).*R_y1_c(y_c+1:Ny+1,x_r);
            bc_t_A2_r[:nh_top, 0] = (C1_y[:nh_top, x_r] * R_y1_c[:nh_top, x_r])
            bc_t_A2_r[Ny + 1 : Ny + y_c + 1, 0] = (C3_y[:nh_top, x_r] * R_y1_c[:nh_top, x_r])
            bc_b_A2_r[y_c : Ny, 0] = (C1_y[y_c + 1 : Ny + 1, x_r] * R_y1_c[y_c + 1 : Ny + 1, x_r])
            bc_b_A2_r[Ny + y_c + 1 : 2 * Ny + 1, 0] = (C3_y[y_c + 1 : Ny + 1, x_r] * R_y1_c[y_c + 1 : Ny + 1, x_r])

        # d_t_xy_A2_r(1:nh_top) = -C2_y(1:nh_top,x_r).*R_y1_c(1:nh_top,x_r);
        # d_b_xy_A2_r(y_c+1:Ny+1) = -C2_y(y_c+1:Ny+1,x_r).*R_y1_c(y_c+1:Ny+1,x_r);
        # d_t_yx_A2_r(2:y_c) = -C4_y(1:nh_top,x_r).*R_y1_c(1:nh_top,x_r);
        # d_b_yx_A2_r(y_c:Ny) = -C4_y(y_c+1:Ny+1,x_r).*R_y1_c(y_c+1:Ny+1,x_r);
            d_t_xy_A2_r[:nh_top, 0] = (-C2_y[:nh_top, x_r] * R_y1_c[:nh_top, x_r])
            d_b_xy_A2_r[y_c + 1: Ny + 1, 0] = (-C2_y[y_c + 1 : Ny + 1, x_r] * R_y1_c[y_c + 1 : Ny + 1, x_r])
            d_t_yx_A2_r[1:y_c + 1, 0] = (-C4_y[:nh_top, x_r] * R_y1_c[:nh_top, x_r])
            d_b_yx_A2_r[y_c : Ny, 0] = (-C4_y[y_c + 1 : Ny + 1, x_r] * R_y1_c[y_c + 1 : Ny + 1, x_r])
            
            A2_l = build_block_diag_A(
                d_c_A2_l.flatten(),
                d_t_A2_l,
                d_b_A2_l,
                bc_t_A2_l.flatten(),
                bc_b_A2_l.flatten(),
                d_t_xy_A2_l.flatten(),
                d_b_xy_A2_l.flatten(),
                d_b_yx_A2_l.flatten(),
                d_t_yx_A2_l.flatten(),
                Ny,
            )
            A2_r = build_block_diag_A(
                d_c_A2_r.flatten(),
                d_t_A2_r,
                d_b_A2_r,
                bc_t_A2_r.flatten(),
                bc_b_A2_r.flatten(),
                d_t_xy_A2_r.flatten(),
                d_b_xy_A2_r.flatten(),
                d_b_yx_A2_r.flatten(),
                d_t_yx_A2_r.flatten(),
                Ny,
            )
            ii = y_c
            bx_2_l[ii, 0] = ( (2 * cx) / (1 + a_x[ii, x_l]) / a_x[ii, x_l] * Wx_m[ii, x_l - 1] + (1 - (2 * cx) / a_x[ii, x_l]) * Wx_m[ii, x_l] * R_in_d[ii, x_l] + (2 * cx) / (1 + a_x[ii, x_l]) * Wx_m[ii, x_l + 1] + R_y1_n[ii, x_l] * C1_y[ii, x_l] * Wx[ii, x_l + 1] + R_y1_n[ii, x_l] * C2_y[ii, x_l] * Wy[ii, x_l + 1])
            by_2_l[ii, 0] = ( (2 * cx) / (1 + a_x[ii, x_l]) / a_x[ii, x_l] * Wy_m[ii, x_l - 1] + (1 - (2 * cx) / a_x[ii, x_l]) * Wy_m[ii, x_l] * R_in_d[ii, x_l] + (2 * cx) / (1 + a_x[ii, x_l]) * Wy_m[ii, x_l + 1] + R_y1_n[ii, x_l] * C3_y[ii, x_l] * Wy[ii, x_l + 1] + R_y1_n[ii, x_l] * C4_y[ii, x_l] * Wx[ii, x_l + 1])
            bx_2_r[ii, 0] = ( (2 * cx) / (1 + a_x[ii, x_r]) * Wx_m[ii, x_r - 1] + (1 - (2 * cx) / a_x[ii, x_r]) * Wx_m[ii, x_r] * R_in_d[ii, x_r] + (2 * cx) / (1 + a_x[ii, x_r]) / a_x[ii, x_r] * Wx_m[ii, x_r + 1] + R_y1_n[ii, x_r] * C1_y[ii, x_r] * Wx[ii, x_r - 1] + R_y1_n[ii, x_r] * C2_y[ii, x_r] * Wy[ii, x_r - 1])
            by_2_r[ii, 0] = ( (2 * cx) / (1 + a_x[ii, x_r]) * Wy_m[ii, x_r - 1] + (1 - (2 * cx) / a_x[ii, x_r]) * Wy_m[ii, x_r] * R_in_d[ii, x_r] + (2 * cx) / (1 + a_x[ii, x_r]) / a_x[ii, x_r] * Wy_m[ii, x_r + 1] + R_y1_n[ii, x_r] * C3_y[ii, x_r] * Wy[ii, x_r - 1] + R_y1_n[ii, x_r] * C4_y[ii, x_r] * Wx[ii, x_r - 1])
            for ii in range(1, min(nh_bottom, nh_top) + 1):
                yy_t = y_c - ii
                yy_b = y_c + ii
                bx_2_l[yy_b, 0] = ( (2 * cx) / (1 + a_x[yy_b, x_l]) / a_x[yy_b, x_l] * Wx_m[yy_b, x_l - 1] + (1 - (2 * cx) / a_x[yy_b, x_l]) * Wx_m[yy_b, x_l] * R_in_d[yy_b, x_l] + (2 * cx) / (1 + a_x[yy_b, x_l]) * Wx_m[yy_b, x_l + 1] + R_y1[yy_b, x_l] * C1_y[yy_b, x_l] * Wx[yy_b - 1, x_l + 1] + R_y1_n[yy_b, x_l] * C1_y[yy_b, x_l] * Wx[yy_b, x_l + 1] + R_y1[yy_b, x_l] * C2_y[yy_b, x_l] * Wy[yy_b - 1, x_l + 1] + R_y1_n[yy_b, x_l] * C2_y[yy_b, x_l] * Wy[yy_b, x_l + 1])
                by_2_l[yy_b, 0] = ( (2 * cx) / (1 + a_x[yy_b, x_l]) / a_x[yy_b, x_l] * Wy_m[yy_b, x_l - 1] + (1 - (2 * cx) / a_x[yy_b, x_l]) * Wy_m[yy_b, x_l] * R_in_d[yy_b, x_l] + (2 * cx) / (1 + a_x[yy_b, x_l]) * Wy_m[yy_b, x_l + 1] + R_y1[yy_b, x_l] * C3_y[yy_b, x_l] * Wy[yy_b - 1, x_l + 1] + R_y1_n[yy_b, x_l] * C3_y[yy_b, x_l] * Wy[yy_b, x_l + 1] + R_y1[yy_b, x_l] * C4_y[yy_b, x_l] * Wx[yy_b - 1, x_l + 1] + R_y1_n[yy_b, x_l] * C4_y[yy_b, x_l] * Wx[yy_b, x_l + 1])
                bx_2_l[yy_t, 0] = ( (2 * cx) / (1 + a_x[yy_t, x_l]) / a_x[yy_t, x_l] * Wx_m[yy_t, x_l - 1] + (1 - (2 * cx) / a_x[yy_t, x_l]) * Wx_m[yy_t, x_l] * R_in_d[yy_t, x_l] + (2 * cx) / (1 + a_x[yy_t, x_l]) * Wx_m[yy_t, x_l + 1] + R_y1[yy_t, x_l] * C1_y[yy_t, x_l] * Wx[yy_t + 1, x_l + 1] + R_y1_n[yy_t, x_l] * C1_y[yy_t, x_l] * Wx[yy_t, x_l + 1] + R_y1[yy_t, x_l] * C2_y[yy_t, x_l] * Wy[yy_t + 1, x_l + 1] + R_y1_n[yy_t, x_l] * C2_y[yy_t, x_l] * Wy[yy_t, x_l + 1])
                by_2_l[yy_t, 0] = ( (2 * cx) / (1 + a_x[yy_t, x_l]) / a_x[yy_t, x_l] * Wy_m[yy_t, x_l - 1] + (1 - (2 * cx) / a_x[yy_t, x_l]) * Wy_m[yy_t, x_l] * R_in_d[yy_t, x_l] + (2 * cx) / (1 + a_x[yy_t, x_l]) * Wy_m[yy_t, x_l + 1] + R_y1[yy_t, x_l] * C3_y[yy_t, x_l] * Wy[yy_t + 1, x_l + 1] + R_y1_n[yy_t, x_l] * C3_y[yy_t, x_l] * Wy[yy_t, x_l + 1] + R_y1[yy_t, x_l] * C4_y[yy_t, x_l] * Wx[yy_t + 1, x_l + 1] + R_y1_n[yy_t, x_l] * C4_y[yy_t, x_l] * Wx[yy_t, x_l + 1])
                bx_2_r[yy_b, 0] = ( (2 * cx) / (1 + a_x[yy_b, x_r]) * Wx_m[yy_b, x_r - 1] + (1 - (2 * cx) / a_x[yy_b, x_r]) * Wx_m[yy_b, x_r] * R_in_d[yy_b, x_r] + (2 * cx) / (1 + a_x[yy_b, x_r]) / a_x[yy_b, x_r] * Wx_m[yy_b, x_r + 1] + R_y1[yy_b, x_r] * C1_y[yy_b, x_r] * Wx[yy_b - 1, x_r - 1] + R_y1_n[yy_b, x_r] * C1_y[yy_b, x_r] * Wx[yy_b, x_r - 1] + R_y1[yy_b, x_r] * C2_y[yy_b, x_r] * Wy[yy_b - 1, x_r - 1] + R_y1_n[yy_b, x_r] * C2_y[yy_b, x_r] * Wy[yy_b, x_r - 1])
                by_2_r[yy_b, 0] = ( (2 * cx) / (1 + a_x[yy_b, x_r]) * Wy_m[yy_b, x_r - 1] + (1 - (2 * cx) / a_x[yy_b, x_r]) * Wy_m[yy_b, x_r] * R_in_d[yy_b, x_r] + (2 * cx) / (1 + a_x[yy_b, x_r]) / a_x[yy_b, x_r] * Wy_m[yy_b, x_r + 1] + R_y1[yy_b, x_r] * C3_y[yy_b, x_r] * Wy[yy_b - 1, x_r - 1] + R_y1_n[yy_b, x_r] * C3_y[yy_b, x_r] * Wy[yy_b, x_r - 1] + R_y1[yy_b, x_r] * C4_y[yy_b, x_r] * Wx[yy_b - 1, x_r - 1] + R_y1_n[yy_b, x_r] * C4_y[yy_b, x_r] * Wx[yy_b, x_r - 1])
                bx_2_r[yy_t, 0] = ( (2 * cx) / (1 + a_x[yy_t, x_r]) * Wx_m[yy_t, x_r - 1] + (1 - (2 * cx) / a_x[yy_t, x_r]) * Wx_m[yy_t, x_r] * R_in_d[yy_t, x_r] + (2 * cx) / (1 + a_x[yy_t, x_r]) / a_x[yy_t, x_r] * Wx_m[yy_t, x_r + 1] + R_y1[yy_t, x_r] * C1_y[yy_t, x_r] * Wx[yy_t + 1, x_r - 1] + R_y1_n[yy_t, x_r] * C1_y[yy_t, x_r] * Wx[yy_t, x_r - 1] + R_y1[yy_t, x_r] * C2_y[yy_t, x_r] * Wy[yy_t + 1, x_r - 1] + R_y1_n[yy_t, x_r] * C2_y[yy_t, x_r] * Wy[yy_t, x_r - 1])
                by_2_r[yy_t, 0] = ( (2 * cx) / (1 + a_x[yy_t, x_r]) * Wy_m[yy_t, x_r - 1] + (1 - (2 * cx) / a_x[yy_t, x_r]) * Wy_m[yy_t, x_r] * R_in_d[yy_t, x_r] + (2 * cx) / (1 + a_x[yy_t, x_r]) / a_x[yy_t, x_r] * Wy_m[yy_t, x_r + 1] + R_y1[yy_t, x_r] * C3_y[yy_t, x_r] * Wy[yy_t + 1, x_r - 1] + R_y1_n[yy_t, x_r] * C3_y[yy_t, x_r] * Wy[yy_t, x_r - 1] + R_y1[yy_t, x_r] * C4_y[yy_t, x_r] * Wx[yy_t + 1, x_r - 1] + R_y1_n[yy_t, x_r] * C4_y[yy_t, x_r] * Wx[yy_t, x_r - 1])
            # 纯上半
            for ii in range(nh_bottom + 1, nh_top + 1):
                yy_t = y_c - ii
                bx_2_l[yy_t, 0] = ( (2 * cx) / (1 + a_x[yy_t, x_l]) / a_x[yy_t, x_l] * Wx_m[yy_t, x_l - 1] + (1 - (2 * cx) / a_x[yy_t, x_l]) * Wx_m[yy_t, x_l] * R_in_d[yy_t, x_l] + (2 * cx) / (1 + a_x[yy_t, x_l]) * Wx_m[yy_t, x_l + 1] + R_y1[yy_t, x_l] * C1_y[yy_t, x_l] * Wx[yy_t + 1, x_l + 1] + R_y1_n[yy_t, x_l] * C1_y[yy_t, x_l] * Wx[yy_t, x_l + 1] + R_y1[yy_t, x_l] * C2_y[yy_t, x_l] * Wy[yy_t + 1, x_l + 1] + R_y1_n[yy_t, x_l] * C2_y[yy_t, x_l] * Wy[yy_t, x_l + 1])
                by_2_l[yy_t, 0] = ( (2 * cx) / (1 + a_x[yy_t, x_l]) / a_x[yy_t, x_l] * Wy_m[yy_t, x_l - 1] + (1 - (2 * cx) / a_x[yy_t, x_l]) * Wy_m[yy_t, x_l] * R_in_d[yy_t, x_l] + (2 * cx) / (1 + a_x[yy_t, x_l]) * Wy_m[yy_t, x_l + 1] + R_y1[yy_t, x_l] * C3_y[yy_t, x_l] * Wy[yy_t + 1, x_l + 1] + R_y1_n[yy_t, x_l] * C3_y[yy_t, x_l] * Wy[yy_t, x_l + 1] + R_y1[yy_t, x_l] * C4_y[yy_t, x_l] * Wx[yy_t + 1, x_l + 1] + R_y1_n[yy_t, x_l] * C4_y[yy_t, x_l] * Wx[yy_t, x_l + 1])
                bx_2_r[yy_t, 0] = ( (2 * cx) / (1 + a_x[yy_t, x_r]) * Wx_m[yy_t, x_r - 1] + (1 - (2 * cx) / a_x[yy_t, x_r]) * Wx_m[yy_t, x_r] * R_in_d[yy_t, x_r] + (2 * cx) / (1 + a_x[yy_t, x_r]) / a_x[yy_t, x_r] * Wx_m[yy_t, x_r + 1] + R_y1[yy_t, x_r] * C1_y[yy_t, x_r] * Wx[yy_t + 1, x_r - 1] + R_y1_n[yy_t, x_r] * C1_y[yy_t, x_r] * Wx[yy_t, x_r - 1] + R_y1[yy_t, x_r] * C2_y[yy_t, x_r] * Wy[yy_t + 1, x_r - 1] + R_y1_n[yy_t, x_r] * C2_y[yy_t, x_r] * Wy[yy_t, x_r - 1])
                by_2_r[yy_t, 0] = ( (2 * cx) / (1 + a_x[yy_t, x_r]) * Wy_m[yy_t, x_r - 1] + (1 - (2 * cx) / a_x[yy_t, x_r]) * Wy_m[yy_t, x_r] * R_in_d[yy_t, x_r] + (2 * cx) / (1 + a_x[yy_t, x_r]) / a_x[yy_t, x_r] * Wy_m[yy_t, x_r + 1] + R_y1[yy_t, x_r] * C3_y[yy_t, x_r] * Wy[yy_t + 1, x_r - 1] + R_y1_n[yy_t, x_r] * C3_y[yy_t, x_r] * Wy[yy_t, x_r - 1] + R_y1[yy_t, x_r] * C4_y[yy_t, x_r] * Wx[yy_t + 1, x_r - 1] + R_y1_n[yy_t, x_r] * C4_y[yy_t, x_r] * Wx[yy_t, x_r - 1])
            # 纯下半
            for ii in range(nh_top + 1, nh_bottom + 1):
                yy_b = y_c + ii
                bx_2_l[yy_b, 0] = ( (2 * cx) / (1 + a_x[yy_b, x_l]) / a_x[yy_b, x_l] * Wx_m[yy_b, x_l - 1] + (1 - (2 * cx) / a_x[yy_b, x_l]) * Wx_m[yy_b, x_l] * R_in_d[yy_b, x_l] + (2 * cx) / (1 + a_x[yy_b, x_l]) * Wx_m[yy_b, x_l + 1] + R_y1[yy_b, x_l] * C1_y[yy_b, x_l] * Wx[yy_b - 1, x_l + 1] + R_y1_n[yy_b, x_l] * C1_y[yy_b, x_l] * Wx[yy_b, x_l + 1] + R_y1[yy_b, x_l] * C2_y[yy_b, x_l] * Wy[yy_b - 1, x_l + 1] + R_y1_n[yy_b, x_l] * C2_y[yy_b, x_l] * Wy[yy_b, x_l + 1])
                by_2_l[yy_b, 0] = ( (2 * cx) / (1 + a_x[yy_b, x_l]) / a_x[yy_b, x_l] * Wy_m[yy_b, x_l - 1] + (1 - (2 * cx) / a_x[yy_b, x_l]) * Wy_m[yy_b, x_l] * R_in_d[yy_b, x_l] + (2 * cx) / (1 + a_x[yy_b, x_l]) * Wy_m[yy_b, x_l + 1] + R_y1[yy_b, x_l] * C3_y[yy_b, x_l] * Wy[yy_b - 1, x_l + 1] + R_y1_n[yy_b, x_l] * C3_y[yy_b, x_l] * Wy[yy_b, x_l + 1] + R_y1[yy_b, x_l] * C4_y[yy_b, x_l] * Wx[yy_b - 1, x_l + 1] + R_y1_n[yy_b, x_l] * C4_y[yy_b, x_l] * Wx[yy_b, x_l + 1])
                bx_2_r[yy_b, 0] = ( (2 * cx) / (1 + a_x[yy_b, x_r]) * Wx_m[yy_b, x_r - 1] + (1 - (2 * cx) / a_x[yy_b, x_r]) * Wx_m[yy_b, x_r] * R_in_d[yy_b, x_r] + (2 * cx) / (1 + a_x[yy_b, x_r]) / a_x[yy_b, x_r] * Wx_m[yy_b, x_r + 1] + R_y1[yy_b, x_r] * C1_y[yy_b, x_r] * Wx[yy_b - 1, x_r - 1] + R_y1_n[yy_b, x_r] * C1_y[yy_b, x_r] * Wx[yy_b, x_r - 1] + R_y1[yy_b, x_r] * C2_y[yy_b, x_r] * Wy[yy_b - 1, x_r - 1] + R_y1_n[yy_b, x_r] * C2_y[yy_b, x_r] * Wy[yy_b, x_r - 1])
                by_2_r[yy_b, 0] = ( (2 * cx) / (1 + a_x[yy_b, x_r]) * Wy_m[yy_b, x_r - 1] + (1 - (2 * cx) / a_x[yy_b, x_r]) * Wy_m[yy_b, x_r] * R_in_d[yy_b, x_r] + (2 * cx) / (1 + a_x[yy_b, x_r]) / a_x[yy_b, x_r] * Wy_m[yy_b, x_r + 1] + R_y1[yy_b, x_r] * C3_y[yy_b, x_r] * Wy[yy_b - 1, x_r - 1] + R_y1_n[yy_b, x_r] * C3_y[yy_b, x_r] * Wy[yy_b, x_r - 1] + R_y1[yy_b, x_r] * C4_y[yy_b, x_r] * Wx[yy_b - 1, x_r - 1] + R_y1_n[yy_b, x_r] * C4_y[yy_b, x_r] * Wx[yy_b, x_r - 1])

            w2_l = np.linalg.solve(
                A2_l, np.concatenate([bx_2_l[:, 0], by_2_l[:, 0]]).astype(complex)
            )
            w2_r = np.linalg.solve(
                A2_r, np.concatenate([bx_2_r[:, 0], by_2_r[:, 0]]).astype(complex)
            )
            Wx[:, x_l] = w2_l[: Ny + 1]
            Wy[:, x_l] = w2_l[Ny + 1 :]
            Wx[:, x_r] = w2_r[: Ny + 1]
            Wy[:, x_r] = w2_r[Ny + 1 :]

        Ex = Wx * np.exp(-1j * k0 * (s_start + z * ds))
        Ey = Wy * np.exp(-1j * k0 * (s_start + z * ds))
        Ex_all[:, :, z] = Ex
        Ey_all[:, :, z] = Ey
        timeCount(f'in z for 6 end')

    return Ex_all, Ey_all
